From 11bb48107f86e3793adb119c7bcc3ec6638e9245 Mon Sep 17 00:00:00 2001
From: Leigh McCulloch <351529+leighmcculloch@users.noreply.github.com>
Date: Thu, 1 Dec 2022 22:10:14 -0800
Subject: [PATCH] Undo breaking api change to registering tokens

---
 soroban-sdk/src/env.rs                 | 37 +++++++++++++--
 soroban-sdk/src/tests.rs               |  1 +
 soroban-sdk/src/tests/contracttoken.rs | 63 ++++++++++++++++++++++++++
 3 files changed, 98 insertions(+), 3 deletions(-)
 create mode 100644 soroban-sdk/src/tests/contracttoken.rs

diff --git a/soroban-sdk/src/env.rs b/soroban-sdk/src/env.rs
index d19853dc0..efbab658c 100644
--- a/soroban-sdk/src/env.rs
+++ b/soroban-sdk/src/env.rs
@@ -588,11 +588,42 @@ impl Env {
     ///
     /// # fn main() {
     /// let env = Env::default();
-    /// env.register_contract_token();
+    /// env.register_contract_token(None);
     /// # }
     /// ```
-    pub fn register_contract_token(&self) -> BytesN<32> {
-        self.register_contract_with_source(xdr::ScContractCode::Token)
+    pub fn register_contract_token<'a>(
+        &self,
+        contract_id: impl Into<Option<&'a BytesN<32>>>,
+    ) -> BytesN<32> {
+        if let Some(contract_id) = contract_id.into() {
+            let contract_id_hash = Hash(contract_id.into());
+            let data_key = xdr::ScVal::Static(xdr::ScStatic::LedgerKeyContractCode);
+            let key = LedgerKey::ContractData(LedgerKeyContractData {
+                contract_id: contract_id_hash.clone(),
+                key: data_key.clone(),
+            });
+            self.env_impl
+                .with_mut_storage(|storage| {
+                    storage.put(
+                        &key,
+                        &LedgerEntry {
+                            ext: xdr::LedgerEntryExt::V0,
+                            last_modified_ledger_seq: 0,
+                            data: xdr::LedgerEntryData::ContractData(xdr::ContractDataEntry {
+                                contract_id: contract_id_hash.clone(),
+                                key: data_key,
+                                val: xdr::ScVal::Object(Some(xdr::ScObject::ContractCode(
+                                    xdr::ScContractCode::Token,
+                                ))),
+                            }),
+                        },
+                    )
+                })
+                .unwrap();
+            contract_id.clone()
+        } else {
+            self.register_contract_with_source(xdr::ScContractCode::Token)
+        }
     }
 
     fn register_contract_with_source(&self, source: xdr::ScContractCode) -> BytesN<32> {
diff --git a/soroban-sdk/src/tests.rs b/soroban-sdk/src/tests.rs
index c98b3ac1f..c8dbae16c 100644
--- a/soroban-sdk/src/tests.rs
+++ b/soroban-sdk/src/tests.rs
@@ -15,3 +15,4 @@ mod contractfile_with_sha256;
 mod contractimport;
 mod contractimport_with_error;
 mod contractimport_with_sha256;
+mod contracttoken;
diff --git a/soroban-sdk/src/tests/contracttoken.rs b/soroban-sdk/src/tests/contracttoken.rs
new file mode 100644
index 000000000..be2d44842
--- /dev/null
+++ b/soroban-sdk/src/tests/contracttoken.rs
@@ -0,0 +1,63 @@
+use crate as soroban_sdk;
+use soroban_sdk::{
+    contractclient, contracttype, testutils::AccountId as _, AccountId, Bytes, BytesN, Env,
+};
+
+#[contractclient(name = "TokenClient")]
+pub trait Token {
+    fn init(env: Env, admin: Identifier, metadata: TokenMetadata);
+    fn name(env: Env) -> Bytes;
+}
+
+#[derive(Clone)]
+#[contracttype]
+pub enum Identifier {
+    Account(AccountId),
+}
+
+#[derive(Clone)]
+#[contracttype]
+pub struct TokenMetadata {
+    pub name: Bytes,
+    pub symbol: Bytes,
+    pub decimals: u32,
+}
+
+#[test]
+fn test_register_token() {
+    let e = Env::default();
+
+    let contract_id = e.register_contract_token(None);
+    let client = TokenClient::new(&e, &contract_id);
+
+    client.init(
+        &Identifier::Account(AccountId::random(&e)),
+        &TokenMetadata {
+            name: Bytes::from_slice(&e, b"testme"),
+            decimals: 7,
+            symbol: Bytes::from_slice(&e, &[]),
+        },
+    );
+
+    assert_eq!(client.name(), Bytes::from_slice(&e, b"testme"))
+}
+
+#[test]
+fn test_register_token_at_id() {
+    let e = Env::default();
+
+    let contract_id = BytesN::from_array(&e, &[1; 32]);
+    e.register_contract_token(&contract_id);
+    let client = TokenClient::new(&e, &contract_id);
+
+    client.init(
+        &Identifier::Account(AccountId::random(&e)),
+        &TokenMetadata {
+            name: Bytes::from_slice(&e, b"testme"),
+            decimals: 7,
+            symbol: Bytes::from_slice(&e, &[]),
+        },
+    );
+
+    assert_eq!(client.name(), Bytes::from_slice(&e, b"testme"))
+}