piecrust: optional owner in migrate

It is now optional to provide an `owner` to the `Session::migrate` call.
If the owner is not specified, it will be taken from the previous
contract at the slot being migrated. If the owner is not specified and
the contract being replaced does not exist, the call will panic.

Resolves #336

piecrust/CHANGELOG.md

@@ -14,6 +14,8 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 ### Changed
 
+- Change `migrate` to take the owner of the contract being replaced if it is
+  not set by the caller [#336]
 - Make `owner` field optional in `ContractData` and `ContractDataBuilder` [#336]
 - Change `ContractData` and `ContractDataBuilder` to take a `Vec<u8>` as owner
   instead of `[u8; N]` [#336]

piecrust/src/session.rs

@@ -406,10 +406,18 @@ impl Session {
     /// given `contract` ID, and the old contract will be removed from the
     /// state.
     ///
+    /// If the `owner` of a contract is not set, it will be set to the owner of
+    /// the contract being replaced. If it is set, then it will be used as the
+    /// new owner.
+    ///
     /// # Errors
     /// The migration may error during execution for a myriad of reasons. The
     /// caller is encouraged to drop the `Session` should an error occur as it
     /// will more than likely be left in an inconsistent state.
+    ///
+    /// # Panics
+    /// If the owner of the new contract is not set in `deploy_data`, and the
+    /// contract being replaced does not exist, this will panic.
     pub fn migrate<'a, A, D, F>(
         mut self,
         contract: ContractId,
@@ -423,8 +431,25 @@ impl Session {
         D: Into<ContractData<'a, A>>,
         F: FnOnce(ContractId, &mut Session) -> Result<(), Error>,
     {
+        let mut new_contract_data = deploy_data.into();
+
+        // If the contract being replaced exists, and the caller did not specify
+        // an owner, set the owner to the owner of the contract being replaced.
+        if let Some(old_contract_data) = self
+            .inner
+            .contract_session
+            .contract(contract)
+            .map_err(|err| PersistenceError(Arc::new(err)))?
+        {
+            if new_contract_data.owner.is_none() {
+                new_contract_data.owner =
+                    Some(old_contract_data.metadata.data().owner.clone());
+            }
+        }
+
         let new_contract =
-            self.deploy(bytecode, deploy_data, deploy_gas_limit)?;
+            self.deploy(bytecode, new_contract_data, deploy_gas_limit)?;
+
         closure(new_contract, &mut self)?;
 
         self.inner

piecrust/tests/persistence.rs

@@ -170,7 +170,7 @@ fn migration() -> Result<(), Error> {
     session = session.migrate(
         contract,
         contract_bytecode!("double_counter"),
-        ContractData::builder().owner(OWNER),
+        ContractData::builder(),
         LIMIT,
         |new_contract, session| {
             let old_counter_value = session
@@ -208,3 +208,72 @@ fn migration() -> Result<(), Error> {
 
     Ok(())
 }
+
+#[test]
+fn migration_new_owner() -> Result<(), Error> {
+    let vm = VM::ephemeral()?;
+    let mut session = vm.session(SessionData::builder())?;
+
+    const OWNER: [u8; 33] = [1u8; 33];
+    const NEW_OWNER: [u8; 33] = [2u8; 33];
+
+    let contract = session.deploy(
+        contract_bytecode!("counter"),
+        ContractData::builder().owner(OWNER),
+        LIMIT,
+    )?;
+
+    let root = session.commit()?;
+
+    let mut session = vm.session(SessionData::builder().base(root))?;
+
+    session = session.migrate(
+        contract,
+        contract_bytecode!("metadata"),
+        ContractData::builder().owner(NEW_OWNER),
+        LIMIT,
+        |_, _| Ok(()),
+    )?;
+
+    let owner = session
+        .call::<_, [u8; 33]>(contract, "read_owner", &(), LIMIT)?
+        .data;
+
+    assert_eq!(owner, NEW_OWNER);
+
+    Ok(())
+}
+
+#[test]
+fn migration_old_owner() -> Result<(), Error> {
+    let vm = VM::ephemeral()?;
+    let mut session = vm.session(SessionData::builder())?;
+
+    const OWNER: [u8; 33] = [1u8; 33];
+
+    let contract = session.deploy(
+        contract_bytecode!("counter"),
+        ContractData::builder().owner(OWNER),
+        LIMIT,
+    )?;
+
+    let root = session.commit()?;
+
+    let mut session = vm.session(SessionData::builder().base(root))?;
+
+    session = session.migrate(
+        contract,
+        contract_bytecode!("metadata"),
+        ContractData::builder(),
+        LIMIT,
+        |_, _| Ok(()),
+    )?;
+
+    let owner = session
+        .call::<_, [u8; 33]>(contract, "read_owner", &(), LIMIT)?
+        .data;
+
+    assert_eq!(owner, OWNER);
+
+    Ok(())
+}