move lock: add writing for managed addresses under env (#16494)

## Description 

Adds a function to write managed addresses for a root package, keyed by
chain ID.

This function is only executed by the test currently, not hooked up to
`sui client publish` or `sui client upgrade` just yet!

See inline comments for more.

## Test Plan 

Added test

---
If your changes are not user-facing and do not break anything, you can
skip the following section. Otherwise, please briefly describe what has
changed under the Release Notes section.

### Type of Change (Check all that apply)

- [ ] protocol change
- [ ] user-visible impact
- [ ] breaking change for a client SDKs
- [ ] breaking change for FNs (FN binary must upgrade)
- [ ] breaking change for validators or node operators (must upgrade
binaries)
- [ ] breaking change for on-chain data layout
- [ ] necessitate either a data wipe or data migration

### Release notes

external-crates/move/crates/move-package/src/lock_file/schema.rs

@@ -7,7 +7,7 @@
 
 use std::io::{Read, Seek, Write};
 
-use anyhow::{bail, Context, Result};
+use anyhow::{anyhow, bail, Context, Result};
 use serde::{Deserialize, Serialize};
 use tempfile::NamedTempFile;
 use toml::value::Value;
@@ -24,6 +24,15 @@ use super::LockFile;
 /// TODO(amnn): Set to version 1 when stabilised.
 pub const VERSION: u64 = 0;
 
+/// Table for storing package info under an environment.
+const ENV_TABLE_NAME: &str = "env";
+
+/// Table keys in environment for managing published packages.
+const ORIGINAL_PUBLISHED_ID_KEY: &str = "original-published-id";
+const LATEST_PUBLISHED_ID_KEY: &str = "latest-published-id";
+const PUBLISHED_VERSION_KEY: &str = "published-version";
+const CHAIN_ID_KEY: &str = "chain-id";
+
 #[derive(Deserialize)]
 pub struct Packages {
     #[serde(rename = "package")]
@@ -237,3 +246,66 @@ fn to_toml_edit_value(value: &toml::Value) -> toml_edit::Item {
         }
     }
 }
+
+pub enum ManagedAddressUpdate {
+    Published {
+        original_id: String,
+        chain_id: String,
+    },
+    Upgraded {
+        latest_id: String,
+        version: u64,
+    },
+}
+
+/// Saves published or upgraded package addresses in the lock file.
+pub fn update_managed_address(
+    file: &mut LockFile,
+    environment: &str,
+    managed_address_update: ManagedAddressUpdate,
+) -> Result<()> {
+    use toml_edit::{value, Document, Item, Table};
+
+    let mut toml_string = String::new();
+    file.read_to_string(&mut toml_string)?;
+    let mut toml = toml_string.parse::<Document>()?;
+
+    let env_table = toml
+        .entry(ENV_TABLE_NAME)
+        .or_insert_with(|| Item::Table(Table::new()))
+        .as_table_mut()
+        .ok_or_else(|| anyhow!("Could not find or create 'env' table in Move.lock"))?
+        .entry(environment)
+        .or_insert_with(|| Item::Table(Table::new()))
+        .as_table_mut()
+        .ok_or_else(|| anyhow!("Could not find or create {environment} table in Move.lock"))?;
+
+    match managed_address_update {
+        ManagedAddressUpdate::Published {
+            original_id,
+            chain_id,
+        } => {
+            env_table[CHAIN_ID_KEY] = value(chain_id);
+            env_table[ORIGINAL_PUBLISHED_ID_KEY] = value(&original_id);
+            env_table[LATEST_PUBLISHED_ID_KEY] = value(original_id);
+            env_table[PUBLISHED_VERSION_KEY] = value("1");
+        }
+        ManagedAddressUpdate::Upgraded { latest_id, version } => {
+            if !env_table.contains_key(CHAIN_ID_KEY) {
+                bail!("Move.lock violation: attempted address update for package upgrade when no {CHAIN_ID_KEY} exists")
+            }
+            if !env_table.contains_key(ORIGINAL_PUBLISHED_ID_KEY) {
+                bail!("Move.lock violation: attempted address update for package upgrade when no {ORIGINAL_PUBLISHED_ID_KEY} exists")
+            }
+            env_table[LATEST_PUBLISHED_ID_KEY] = value(latest_id);
+            env_table[PUBLISHED_VERSION_KEY] = value(version.to_string());
+        }
+    }
+
+    file.set_len(0)?;
+    file.rewind()?;
+    write!(file, "{}", toml)?;
+    file.flush()?;
+    file.rewind()?;
+    Ok(())
+}

external-crates/move/crates/move-package/tests/test_lock_file.rs

@@ -10,7 +10,9 @@ use std::{
 use tempfile::TempDir;
 
 use move_compiler::editions::{Edition, Flavor};
-use move_package::lock_file::schema::ToolchainVersion;
+use move_package::lock_file::schema::{
+    update_managed_address, ManagedAddressUpdate, ToolchainVersion,
+};
 use move_package::lock_file::LockFile;
 use move_package::BuildConfig;
 
@@ -101,6 +103,71 @@ fn update_lock_file_toolchain_version() {
     expected.assert_eq(&toml);
 }
 
+#[test]
+fn test_update_managed_address() {
+    let pkg = create_test_package().unwrap();
+    let lock_path = pkg.path().join("Move.lock");
+
+    // Initialize lock file.
+    let lock = LockFile::new(
+        pkg.path().to_path_buf(),
+        /* manifest_digest */ "42".to_string(),
+        /* deps_digest */ "7".to_string(),
+    )
+    .unwrap();
+    lock.commit(&lock_path).unwrap();
+
+    // Update managed address in lock file.
+    let pb = PathBuf::from(pkg.path());
+    let mut lock = LockFile::from(pb, &lock_path).unwrap();
+    update_managed_address(
+        &mut lock,
+        "default".into(),
+        ManagedAddressUpdate::Published {
+            original_id: "0x123".into(),
+            chain_id: "35834a8a".into(),
+        },
+    )
+    .unwrap();
+
+    update_managed_address(
+        &mut lock,
+        "default".into(),
+        ManagedAddressUpdate::Upgraded {
+            latest_id: "0x456".into(),
+            version: 2,
+        },
+    )
+    .unwrap();
+    lock.commit(&lock_path).unwrap();
+
+    // Read lock file and check contents.
+    let mut lock_file = File::open(lock_path).unwrap();
+    let contents = {
+        let mut buf = String::new();
+        let _ = lock_file.read_to_string(&mut buf);
+        buf
+    };
+
+    let expected = expect![[r##"
+        # @generated by Move, please check-in and do not edit manually.
+
+        [move]
+        version = 0
+        manifest_digest = "42"
+        deps_digest = "7"
+
+        [env]
+
+        [env.default]
+        chain-id = "35834a8a"
+        original-published-id = "0x123"
+        latest-published-id = "0x456"
+        published-version = "2"
+    "##]];
+    expected.assert_eq(&contents);
+}
+
 /// Create a simple Move package with no sources (just a manifest and an output directory) in a
 /// temporary directory, and return it.
 fn create_test_package() -> io::Result<TempDir> {