Upgradable contracts using set_code function

frame/contracts/fixtures/new_set_code_hash_contract.wat

@@ -0,0 +1,13 @@
+(module
+	(import "seal0" "seal_return" (func $seal_return (param i32 i32 i32)))
+	(import "env" "memory" (memory 1 1))
+
+	;; [0, 32) return value
+	(data (i32.const 0) "\02")
+
+	(func (export "deploy"))
+
+	(func (export "call")
+		(call $seal_return (i32.const 0) (i32.const 0) (i32.const 4))
+	)
+)

frame/contracts/fixtures/set_code_hash.wat

@@ -0,0 +1,43 @@
+(module
+	(import "seal0" "seal_input" (func $seal_input (param i32 i32)))
+	(import "seal0" "seal_return" (func $seal_return (param i32 i32 i32)))
+	(import "__unstable__" "seal_set_code_hash" (func $seal_set_code_hash (param i32) (result i32)))
+
+	(import "env" "memory" (memory 1 1))
+
+	;; [0, 32) here we store input
+
+	;; [32, 36) input size
+	(data (i32.const 32) "\20")
+
+	;; [36, 40) return value
+	(data (i32.const 36) "\01")
+
+	(func $assert (param i32)
+		(block $ok
+			(br_if $ok
+				(get_local 0)
+			)
+			(unreachable)
+		)
+	)
+
+	(func (export "call")
+		(local $exit_code i32)
+
+		(call $seal_input (i32.const 0) (i32.const 32))
+
+		(set_local $exit_code
+			(call $seal_set_code_hash (i32.const 0)) ;; Pointer to the input data.
+		)
+		(call $assert
+			(i32.eq (get_local $exit_code) (i32.const 0)) ;; ReturnCode::Success
+		)
+
+		;; we return 1 after setting new code_hash
+		;; next `call` will NOT return this value, because contract code has been changed
+		(call $seal_return (i32.const 0) (i32.const 36) (i32.const 4))
+	)
+
+	(func (export "deploy"))
+)

frame/contracts/src/benchmarking/mod.rs

@@ -808,6 +808,47 @@ benchmarks! {
 		let origin = RawOrigin::Signed(instance.caller.clone());
 	}: call(origin, instance.addr, 0u32.into(), Weight::MAX, None, vec![])
 
+	#[skip_meta]
+	seal_set_code_hash {
+		let r in 0 .. API_BENCHMARK_BATCHES;
+		let code_hashes = (0..r * API_BENCHMARK_BATCH_SIZE)
+			.map(|i| {
+				let new_code = WasmModule::<T>::dummy_with_bytes(i);
+				Contracts::<T>::store_code_raw(new_code.code, whitelisted_caller())?;
+				Ok(new_code.hash)
+			})
+			.collect::<Result<Vec<_>, &'static str>>()?;
+		let code_hash_len = code_hashes.get(0).map(|x| x.encode().len()).unwrap_or(0);
+		let code_hashes_bytes = code_hashes.iter().flat_map(|x| x.encode()).collect::<Vec<_>>();
+		let code_hashes_len = code_hashes_bytes.len();
+
+		let code = WasmModule::<T>::from(ModuleDefinition {
+			memory: Some(ImportedMemory::max::<T>()),
+			imported_functions: vec![ImportedFunction {
+				module: "__unstable__",
+				name: "seal_set_code_hash",
+				params: vec![
+					ValueType::I32,
+				],
+				return_type: Some(ValueType::I32),
+			}],
+			data_segments: vec![
+				DataSegment {
+					offset: 0 as u32,
+					value: code_hashes_bytes,
+				},
+			],
+			call_body: Some(body::repeated_dyn(r * API_BENCHMARK_BATCH_SIZE, vec![
+				Counter(0 as u32, code_hash_len as u32), // code_hash_ptr
+				Regular(Instruction::Call(0)),
+				Regular(Instruction::Drop),
+			])),
+			.. Default::default()
+		});
+		let instance = Contract::<T>::new(code, vec![])?;
+		let origin = RawOrigin::Signed(instance.caller.clone());
+	}: call(origin, instance.addr, 0u32.into(), Weight::MAX, None, vec![])
+
 	seal_set_storage_per_kb {
 		let n in 0 .. T::Schedule::get().limits.payload_len / 1024;
 		let key = T::Hashing::hash_of(&1u32).as_ref().to_vec();

frame/contracts/src/exec.rs

@@ -18,6 +18,7 @@
 use crate::{
 	gas::GasMeter,
 	storage::{self, Storage, WriteOutcome},
+	wasm::{decrement_refcount, increment_refcount},
 	AccountCounter, BalanceOf, CodeHash, Config, ContractInfo, ContractInfoOf, Error, Event,
 	Pallet as Contracts, Schedule,
 };
@@ -155,6 +156,9 @@ pub trait Ext: sealing::Sealed {
 		take_old: bool,
 	) -> Result<WriteOutcome, DispatchError>;
 
+	/// Sets new code hash for existing contract
+	fn set_code_hash(&mut self, hash: CodeHash<Self::T>) -> Result<(), DispatchError>;
+
 	/// Returns a reference to the account id of the caller.
 	fn caller(&self) -> &AccountIdOf<Self::T>;
 
@@ -1016,6 +1020,19 @@ where
 		)
 	}
 
+	fn set_code_hash(&mut self, hash: CodeHash<Self::T>) -> Result<(), DispatchError> {
+		let top_frame = &mut self.top_frame_mut();
+		let prev_hash = top_frame.contract_info().code_hash.clone();
+		top_frame.contract_info().code_hash = hash;
+		increment_refcount::<Self::T>(hash)?;
+		decrement_refcount::<Self::T>(prev_hash)?;
+		Contracts::<Self::T>::deposit_event(Event::ContractCodeUpdated {
+			contract: top_frame.account_id.clone(),
+			code_hash: hash,
+		});
+		Ok(())
+	}
+
 	fn address(&self) -> &T::AccountId {
 		&self.top_frame().account_id
 	}

frame/contracts/src/lib.rs

@@ -559,6 +559,14 @@ pub mod pallet {
 
 		/// A code with the specified hash was removed.
 		CodeRemoved { code_hash: T::Hash },
+
+		/// Contract code has been updated to one specified with code_hash
+		ContractCodeUpdated {
+			/// The contract that has been updated.
+			contract: T::AccountId,
+			/// New code hash that was set for the contract
+			code_hash: T::Hash,
+		},
 	}
 
 	#[pallet::error]

frame/contracts/src/schedule.rs

@@ -319,6 +319,9 @@ pub struct HostFnWeights<T: Config> {
 	/// Weight per byte of an item stored with `seal_set_storage`.
 	pub set_storage_per_byte: Weight,
 
+	/// Weight of calling `seal_set_code_hash`.
+	pub set_code_hash: Weight,
+
 	/// Weight of calling `seal_clear_storage`.
 	pub clear_storage: Weight,
 
@@ -587,6 +590,7 @@ impl<T: Config> Default for HostFnWeights<T> {
 			debug_message: cost_batched!(seal_debug_message),
 			set_storage: cost_batched!(seal_set_storage),
 			set_storage_per_byte: cost_byte_batched!(seal_set_storage_per_kb),
+			set_code_hash: cost_batched!(seal_set_code_hash),
 			clear_storage: cost_batched!(seal_clear_storage),
 			contains_storage: cost_batched!(seal_contains_storage),
 			get_storage: cost_batched!(seal_get_storage),

frame/contracts/src/tests.rs

@@ -691,6 +691,66 @@ fn deploy_and_call_other_contract() {
 	});
 }
 
+#[test]
+#[cfg(feature = "unstable-interface")]
+fn set_code_hash() {
+	let (wasm, code_hash) = compile_module::<Test>("set_code_hash").unwrap();
+	let (new_wasm, new_code_hash) = compile_module::<Test>("new_set_code_hash_contract").unwrap();
+
+	let contract_addr = Contracts::contract_address(&ALICE, &code_hash, &[]);
+
+	ExtBuilder::default().existential_deposit(100).build().execute_with(|| {
+		let _ = Balances::deposit_creating(&ALICE, 1_000_000);
+
+		// Instantiate the 'caller'
+		assert_ok!(Contracts::instantiate_with_code(
+			Origin::signed(ALICE),
+			300_000,
+			GAS_LIMIT,
+			None,
+			wasm,
+			vec![],
+			vec![],
+		));
+		// upload new code
+		assert_ok!(Contracts::upload_code(Origin::signed(ALICE), new_wasm.clone(), None));
+
+		// First call sets new code_hash and returns 1
+		let result = Contracts::bare_call(
+			ALICE,
+			contract_addr.clone(),
+			0,
+			GAS_LIMIT,
+			None,
+			new_code_hash.as_ref().to_vec(),
+			true,
+		)
+		.result
+		.unwrap();
+		assert_return_code!(result, 1);
+
+		// Second calls new contract code that returns 2
+		let result =
+			Contracts::bare_call(ALICE, contract_addr.clone(), 0, GAS_LIMIT, None, vec![], true)
+				.result
+				.unwrap();
+		assert_return_code!(result, 2);
+
+		// Checking for the last event only
+		assert_eq!(
+			System::events()[13],
+			EventRecord {
+				phase: Phase::Initialization,
+				event: Event::Contracts(crate::Event::ContractCodeUpdated {
+					contract: contract_addr.clone(),
+					code_hash: new_code_hash.clone(),
+				}),
+				topics: vec![],
+			},
+		);
+	});
+}
+
 #[test]
 fn cannot_self_destruct_through_draning() {
 	let (wasm, code_hash) = compile_module::<Test>("drain").unwrap();

frame/contracts/src/wasm/code_cache.rs

@@ -118,6 +118,22 @@ pub fn decrement_refcount<T: Config>(code_hash: CodeHash<T>) -> Result<(), Dispa
 	Ok(())
 }
 
+/// Increment the refcount of a code in-storage by one.
+///
+/// # Errors
+///
+/// DispatchError::CannotLookup is returned if there is no info for specified code_hash
+pub fn increment_refcount<T: Config>(code_hash: CodeHash<T>) -> Result<(), DispatchError> {
+	<OwnerInfoOf<T>>::mutate(code_hash, |existing| -> Result<(), DispatchError> {
+		if let Some(info) = existing {
+			info.refcount = info.refcount.saturating_add(1);
+			Ok(())
+		} else {
+			Err(Error::<T>::CodeNotFound.into())
+		}
+	})
+}
+
 /// Try to remove code together with all associated information.
 pub fn try_remove<T: Config>(origin: &T::AccountId, code_hash: CodeHash<T>) -> DispatchResult {
 	<OwnerInfoOf<T>>::try_mutate_exists(&code_hash, |existing| {

frame/contracts/src/wasm/mod.rs

@@ -26,7 +26,10 @@ mod runtime;
 
 #[cfg(feature = "runtime-benchmarks")]
 pub use self::code_cache::reinstrument;
-pub use self::runtime::{ReturnCode, Runtime, RuntimeCosts};
+pub use self::{
+	code_cache::{decrement_refcount, increment_refcount},
+	runtime::{ReturnCode, Runtime, RuntimeCosts},
+};
 use crate::{
 	exec::{ExecResult, Executable, ExportedFunction, Ext},
 	gas::GasMeter,
@@ -303,6 +306,7 @@ mod tests {
 	}
 
 	pub struct MockExt {
+		code_hashes: Vec<CodeHash<Test>>,
 		storage: HashMap<StorageKey, Vec<u8>>,
 		instantiates: Vec<InstantiateEntry>,
 		terminations: Vec<TerminationEntry>,
@@ -325,6 +329,7 @@ mod tests {
 	impl Default for MockExt {
 		fn default() -> Self {
 			Self {
+				code_hashes: Default::default(),
 				storage: Default::default(),
 				instantiates: Default::default(),
 				terminations: Default::default(),
@@ -406,6 +411,10 @@ mod tests {
 			}
 			Ok(result)
 		}
+		fn set_code_hash(&mut self, hash: CodeHash<Self::T>) -> Result<(), DispatchError> {
+			self.code_hashes.push(hash);
+			Ok(())
+		}
 		fn caller(&self) -> &AccountIdOf<Self::T> {
 			&ALICE
 		}
@@ -2180,6 +2189,47 @@ mod tests {
 		assert_eq!(&result.data.0[4..], &[0u8; 0]);
 	}
 
+	#[test]
+	#[cfg(feature = "unstable-interface")]
+	fn set_code_hash() {
+		const CODE: &str = r#"
+(module
+	(import "__unstable__" "seal_set_code_hash" (func $seal_set_code_hash (param i32) (result i32)))
+	(import "env" "memory" (memory 1 1))
+	(func $assert (param i32)
+		(block $ok
+			(br_if $ok
+				(get_local 0)
+			)
+			(unreachable)
+		)
+	)
+	(func (export "call")
+		(local $exit_code i32)
+		(set_local $exit_code
+			(call $seal_set_code_hash (i32.const 0))
+		)
+		(call $assert
+			(i32.eq (get_local $exit_code) (i32.const 0)) ;; ReturnCode::Success
+		)
+	)
+
+	(func (export "deploy"))
+
+	;; Hash of code.
+	(data (i32.const 0)
+		"\11\11\11\11\11\11\11\11\11\11\11\11\11\11\11\11"
+		"\11\11\11\11\11\11\11\11\11\11\11\11\11\11\11\11"
+	)
+)
+"#;
+
+		let mut mock_ext = MockExt::default();
+		execute(CODE, [0u8; 32].encode(), &mut mock_ext).unwrap();
+
+		assert_eq!(mock_ext.code_hashes.pop().unwrap(), H256::from_slice(&[17u8; 32]));
+	}
+
 	#[test]
 	#[cfg(feature = "unstable-interface")]
 	fn contains_storage_works() {