From bbb1354f50c665fa794d125491edf3904637f2cb Mon Sep 17 00:00:00 2001
From: Milosz Muszynski <milosz.muszynski@gmail.com>
Date: Wed, 24 Jan 2024 16:56:46 +0100
Subject: [PATCH] rusk-abi: introduced raw owner calls

---
 rusk-abi/src/abi.rs                         | 34 +++++++++++++++++++--
 rusk-abi/src/lib.rs                         |  3 --
 rusk-abi/src/types.rs                       |  4 ---
 rusk-abi/tests/contracts/host_fn/src/lib.rs |  2 +-
 4 files changed, 32 insertions(+), 11 deletions(-)

diff --git a/rusk-abi/src/abi.rs b/rusk-abi/src/abi.rs
index af5e828696..7c7a01b991 100644
--- a/rusk-abi/src/abi.rs
+++ b/rusk-abi/src/abi.rs
@@ -5,7 +5,12 @@
 // Copyright (c) DUSK NETWORK. All rights reserved.
 
 #[cfg(feature = "abi")]
-use crate::OwnerRaw;
+use dusk_bytes::Serializable;
+#[cfg(feature = "abi")]
+use dusk_pki::PublicSpendKey;
+#[cfg(feature = "abi")]
+use rkyv::{archived_root, Deserialize, Infallible};
+// re-export `piecrust-uplink` such that `rusk-abi` is the only crate
 pub use piecrust_uplink::*;
 
 /// Compute the blake2b hash of the given bytes, returning the resulting scalar.
@@ -75,6 +80,29 @@ pub fn payment_info(
 
 /// Query owner of a given contract.
 #[cfg(feature = "abi")]
-pub fn owner_raw(contract: ContractId) -> Option<OwnerRaw> {
-    owner(contract)
+pub fn owner(contract: ContractId) -> Option<PublicSpendKey> {
+    owner_raw(contract).map(|buf| {
+        let ret = unsafe { archived_root::<PublicSpendKey>(buf.as_slice()) };
+        ret.deserialize(&mut Infallible).expect("Infallible")
+    })
+}
+
+/// Query owner of a given contract.
+#[cfg(feature = "abi")]
+pub fn self_owner() -> PublicSpendKey {
+    let buf = self_owner_raw();
+    let ret = unsafe { archived_root::<PublicSpendKey>(buf.as_slice()) };
+    ret.deserialize(&mut Infallible).expect("Infallible")
+}
+
+/// Query raw owner of a given contract.
+#[cfg(feature = "abi")]
+pub fn owner_raw(contract: ContractId) -> Option<[u8; PublicSpendKey::SIZE]> {
+    piecrust_uplink::owner(contract)
+}
+
+/// Query raw self owner.
+#[cfg(feature = "abi")]
+pub fn self_owner_raw() -> [u8; PublicSpendKey::SIZE] {
+    piecrust_uplink::self_owner()
 }
diff --git a/rusk-abi/src/lib.rs b/rusk-abi/src/lib.rs
index 33371407f5..6df9abffc8 100644
--- a/rusk-abi/src/lib.rs
+++ b/rusk-abi/src/lib.rs
@@ -39,9 +39,6 @@ pub mod hash;
 
 use hash::Hasher;
 
-// re-export `piecrust-uplink` such that `rusk-abi` is the only crate
-pub use piecrust_uplink::*;
-
 use dusk_bls12_381::BlsScalar;
 use dusk_bytes::DeserializableSlice;
 
diff --git a/rusk-abi/src/types.rs b/rusk-abi/src/types.rs
index 0126451dca..9e28496805 100644
--- a/rusk-abi/src/types.rs
+++ b/rusk-abi/src/types.rs
@@ -11,7 +11,6 @@ use dusk_jubjub::{JubJubAffine, JubJubExtended, JubJubScalar};
 use dusk_pki::PublicSpendKey;
 
 use bytecheck::CheckBytes;
-use dusk_bytes::Serializable;
 use rkyv::{Archive, Deserialize, Serialize};
 
 pub(crate) enum Query {}
@@ -93,6 +92,3 @@ where
         t.clone().into()
     }
 }
-
-/// Raw form of an owner type
-pub type OwnerRaw = [u8; PublicSpendKey::SIZE];
diff --git a/rusk-abi/tests/contracts/host_fn/src/lib.rs b/rusk-abi/tests/contracts/host_fn/src/lib.rs
index a2f11615ab..a0c636fe49 100644
--- a/rusk-abi/tests/contracts/host_fn/src/lib.rs
+++ b/rusk-abi/tests/contracts/host_fn/src/lib.rs
@@ -69,7 +69,7 @@ impl HostFnTest {
     }
 
     pub fn owner(&self) -> [u8; PublicSpendKey::SIZE] {
-        rusk_abi::self_owner()
+        rusk_abi::self_owner_raw()
     }
 }