Use customized ink address as contract id

crates/phactory/src/contracts/pink.rs

@@ -64,7 +64,7 @@ impl Pink {
         Self { instance, group }
     }
 
-    pub fn address_to_id(address: &AccountId) -> contracts::NativeContractId {
+    pub fn address_to_id(address: &AccountId) -> contracts::ContractId {
         let inner: &[u8; 32] = address.as_ref();
         inner.into()
     }
@@ -161,7 +161,7 @@ impl contracts::NativeContract for Pink {
 }
 
 impl NativeContractMore for Pink {
-    fn id(&self) -> contracts::NativeContractId {
+    fn id(&self) -> phala_mq::ContractId {
         Pink::address_to_id(&self.instance.address)
     }
 
@@ -208,8 +208,7 @@ pub mod group {
             block_number: BlockNumber,
             now: u64,
         ) -> Result<ExecSideEffects> {
-            let group = self
-                .get_group_or_default_mut(&group_id, contract_key);
+            let group = self.get_group_or_default_mut(&group_id, contract_key);
             let (_, effects) = Pink::instantiate(
                 group_id,
                 &mut group.storage,
@@ -239,13 +238,15 @@ pub mod group {
             group_id: &ContractGroupId,
             contract_key: &sr25519::Pair,
         ) -> &mut Group {
-            self.groups
-                .entry(group_id.clone())
-                .or_insert_with(|| Group {
+            self.groups.entry(group_id.clone()).or_insert_with(|| {
+                let mut group = Group {
                     storage: Default::default(),
                     contracts: Default::default(),
                     key: contract_key.clone(),
-                })
+                };
+                group.set_id(group_id);
+                group
+            })
         }
 
         pub fn commit_changes(&mut self) -> anyhow::Result<()> {
@@ -278,5 +279,9 @@ pub mod group {
             self.storage.commit_changes();
             Ok(())
         }
+
+        pub fn set_id(&mut self, id: &ContractGroupId) {
+            self.storage.set_group_id(id.as_bytes());
+        }
     }
 }

crates/phactory/src/contracts/support.rs

@@ -54,23 +54,8 @@ pub trait Contract {
     fn set_on_block_end_selector(&mut self, selector: u32);
 }
 
-#[derive(Encode, Decode, Clone, Debug, Copy, derive_more::From)]
-pub struct NativeContractId(sp_core::H256);
-
-impl From<&[u8; 32]> for NativeContractId {
-    fn from(bytes: &[u8; 32]) -> Self {
-        NativeContractId(bytes.into())
-    }
-}
-
-impl NativeContractId {
-    pub fn to_contract_id(&self, group_id: &phala_mq::ContractGroupId) -> phala_mq::ContractId {
-        sp_core::blake2_256(&(group_id, &self.0).encode()).into()
-    }
-}
-
 pub trait NativeContractMore {
-    fn id(&self) -> NativeContractId;
+    fn id(&self) -> phala_mq::ContractId;
     fn set_on_block_end_selector(&mut self, _selector: u32) {}
 }
 
@@ -105,8 +90,14 @@ pub struct NativeContractWrapper<Con> {
 }
 
 impl<Con> NativeContractWrapper<Con> {
-    pub fn new(inner: Con, deployer: sp_core::H256, salt: &[u8], id: u32) -> Self {
-        let encoded = (deployer, id, salt).encode();
+    pub fn new(
+        inner: Con,
+        group_id: &phala_mq::ContractGroupId,
+        deployer: sp_core::H256,
+        salt: &[u8],
+        id: u32,
+    ) -> Self {
+        let encoded = (deployer, id, group_id, salt).encode();
         let id = sp_core::blake2_256(&encoded).into();
         NativeContractWrapper { inner, id }
     }
@@ -141,8 +132,8 @@ impl<Con: NativeContract> NativeContract for NativeContractWrapper<Con> {
 }
 
 impl<Con: NativeContract> NativeContractMore for NativeContractWrapper<Con> {
-    fn id(&self) -> NativeContractId {
-        self.id.into()
+    fn id(&self) -> phala_mq::ContractId {
+        self.id
     }
 }
 

crates/phactory/src/system/mod.rs

@@ -864,6 +864,7 @@ impl<Platform: pal::Platform> System<Platform> {
                                 $id => {
                                     let contract = NativeContractWrapper::new(
                                         $contract,
+                                        &group_id,
                                         deployer,
                                         &salt,
                                         $id,
@@ -1073,7 +1074,7 @@ pub fn apply_pink_side_effects(
     }
 
     for (address, event) in effects.pink_events {
-        let id = Pink::address_to_id(&address).to_contract_id(&group_id);
+        let id = Pink::address_to_id(&address);
         let contract = match contracts.get_mut(&id) {
             Some(contract) => contract,
             None => {
@@ -1116,7 +1117,7 @@ where
     <Contract as NativeContract>::Cmd: Send,
     contracts::AnyContract: From<contracts::NativeCompatContract<Contract>>,
 {
-    let contract_id = contract.id().to_contract_id(&group_id);
+    let contract_id = contract.id();
     if contracts.get(&contract_id).is_some() {
         return Err(anyhow::anyhow!("Contract already exists"));
     }

crates/pink/src/runtime.rs

@@ -1,5 +1,6 @@
 mod extension;
 mod mock_types;
+mod pallet_pink;
 
 use crate::types::{AccountId, Balance, BlockNumber, Hash, Hashing, Index};
 use frame_support::weights::Weight;
@@ -27,6 +28,7 @@ frame_support::construct_runtime! {
         Timestamp: pallet_timestamp::{Pallet, Call, Storage, Inherent},
         Randomness: pallet_randomness_collective_flip::{Pallet, Storage},
         Contracts: pallet_contracts::{Pallet, Call, Storage, Event<T>},
+        Pink: pallet_pink::{Pallet, Storage},
     }
 }
 
@@ -37,6 +39,8 @@ parameter_types! {
     pub static ExistentialDeposit: u64 = 0;
 }
 
+impl pallet_pink::Config for PinkRuntime {}
+
 impl frame_system::Config for PinkRuntime {
     type BaseCallFilter = frame_support::traits::Everything;
     type BlockWeights = BlockWeights;
@@ -113,6 +117,7 @@ impl Config for PinkRuntime {
     type DeletionWeightLimit = DeletionWeightLimit;
     type Schedule = DefaultSchedule;
     type ContractDeposit = ();
+    type AddressGenerator = Pink;
 }
 
 #[cfg(test)]
@@ -212,10 +217,9 @@ mod tests {
                 // We offset data in the contract tables by 1.
                 let mut params = vec![(n + 1) as u8];
                 params.extend_from_slice(input);
-                let result =
-                    Contracts::bare_call(ALICE, addr.clone(), 0, GAS_LIMIT, params, false)
-                        .result
-                        .unwrap();
+                let result = Contracts::bare_call(ALICE, addr.clone(), 0, GAS_LIMIT, params, false)
+                    .result
+                    .unwrap();
                 assert!(result.is_success());
                 let expected = hash_fn(input.as_ref());
                 assert_eq!(&result.data[..*expected_size], &*expected);
@@ -225,9 +229,7 @@ mod tests {
 
     pub mod exec {
         use sp_runtime::traits::BlakeTwo256;
-        use sp_state_machine::{
-            Backend, Ext, OverlayedChanges, StorageTransactionCache,
-        };
+        use sp_state_machine::{Backend, Ext, OverlayedChanges, StorageTransactionCache};
         pub type InMemoryBackend = sp_state_machine::InMemoryBackend<BlakeTwo256>;
 
         pub fn execute_with<R>(f: impl FnOnce() -> R) -> R {

crates/pink/src/runtime/pallet_pink.rs

@@ -0,0 +1,50 @@
+pub use pallet::*;
+
+#[frame_support::pallet]
+pub mod pallet {
+    use frame_support::pallet_prelude::*;
+    use pallet_contracts::AddressGenerator;
+    use sp_core::crypto::UncheckedFrom;
+    use sp_runtime::traits::Hash as _;
+
+    type CodeHash<T> = <T as frame_system::Config>::Hash;
+
+    #[pallet::config]
+    pub trait Config: frame_system::Config {}
+
+    #[pallet::storage]
+    pub(crate) type GroupId<T: Config> = StorageValue<_, Vec<u8>, ValueQuery>;
+
+    #[pallet::pallet]
+    pub struct Pallet<T>(PhantomData<T>);
+
+    impl<T: Config> AddressGenerator<T> for Pallet<T>
+    where
+        T::AccountId: UncheckedFrom<T::Hash> + AsRef<[u8]>,
+    {
+        fn generate_address(
+            deploying_address: &T::AccountId,
+            code_hash: &CodeHash<T>,
+            salt: &[u8],
+        ) -> T::AccountId {
+            let group_id = <GroupId<T>>::get();
+            let buf: Vec<_> = deploying_address
+                .as_ref()
+                .iter()
+                .chain(code_hash.as_ref())
+                .chain(&group_id)
+                .chain(salt)
+                .cloned()
+                .collect();
+            UncheckedFrom::unchecked_from(<T as frame_system::Config>::Hashing::hash(
+                &buf,
+            ))
+        }
+    }
+
+    impl<T: Config> Pallet<T> {
+        pub fn set_group_id(group_id: &[u8]) {
+            <GroupId<T>>::put(group_id.to_vec());
+        }
+    }
+}

crates/pink/src/storage.rs

@@ -87,6 +87,12 @@ where
         self.commit_transaction(root, transaction);
         self.clear_changes();
     }
+
+    pub fn set_group_id(&mut self, group_id: &[u8]) {
+        self.execute_with(false, || {
+            crate::runtime::Pink::set_group_id(group_id);
+        });
+    }
 }
 
 impl Serialize for Storage<InMemoryBackend> {