add: CPC extension for SelfCall

cli/src/runner.rs

@@ -70,7 +70,7 @@ pub fn tapes_to_runtime_arguments(
         .call_tape
         .writer
         .iter()
-        .map(|msg| msg.callee)
+        .map(|msg| msg.callee.0)
         .collect::<BTreeSet<_>>()
         .into_iter()
         .collect_vec();

examples/Cargo.toml

@@ -16,6 +16,7 @@ members = [
   "token",
   "wallet",
   "inputtape",
+  "selfcaller",
 ]
 resolver = "2"
 

examples/inputtape/main_native.rs

@@ -2,7 +2,9 @@
 #![allow(unused_attributes)]
 mod core_logic;
 
-use mozak_sdk::common::types::{Poseidon2Hash, ProgramIdentifier};
+use mozak_sdk::common::types::{
+    Poseidon2Hash, ProgramIdentifier, SelfCallExtendedProgramIdentifier, SelfCallExtensionFlag,
+};
 
 use crate::core_logic::{dispatch, MethodArgs};
 
@@ -12,7 +14,10 @@ fn main() {
     let buf1 = Poseidon2Hash::new_from_rand_seed(2).inner();
     let buf2 = buf1.iter().map(|x| x.wrapping_add(1)).collect::<Vec<u8>>();
 
-    mozak_sdk::add_identity(token_program); // Manual override for `IdentityStack`
+    mozak_sdk::add_identity(SelfCallExtendedProgramIdentifier(
+        token_program,
+        SelfCallExtensionFlag::default(),
+    )); // Manual override for `IdentityStack`
     let _ = mozak_sdk::write(&mozak_sdk::InputTapeType::PublicTape, &buf1);
     let _ = mozak_sdk::write(&mozak_sdk::InputTapeType::PrivateTape, &buf2[..]);
     mozak_sdk::rm_identity(); // Manual override for `IdentityStack`

examples/selfcaller/Cargo.toml

@@ -0,0 +1,34 @@
+[package]
+edition = "2021"
+name = "selfcaller"
+version = "0.1.0"
+
+[features]
+native = []
+
+[dependencies]
+mozak-sdk = { path = "../../sdk" }
+rkyv = { version = "=0.8.0-alpha.1", default-features = false, features = ["pointer_width_32", "alloc"] }
+rkyv_derive = "=0.8.0-alpha.1"
+
+[target.'cfg(not(target_os="mozakvm"))'.dependencies]
+hex = "0.4"
+rand = "0.8"
+rand_chacha = "0.3"
+
+[[bin]]
+name = "selfcallerbin"
+path = "main_mozak.rs"
+
+[[bin]]
+name = "selfcaller-native"
+path = "main_native.rs"
+required-features = ["native"]
+
+[lib]
+name = "selfcaller"
+path = "core_logic.rs"
+
+# The following is read by `run_examples.py`
+[package.metadata.mozak]
+example_program_id = "MZK-5b7b6135be198533f7c7ec46651216b762e6d47e69b408d1bc79d641f9ae06de"

examples/selfcaller/README.md

@@ -0,0 +1,34 @@
+# To run
+
+## Native
+
+To run on your system, use the following command (kindly change [target triple](https://doc.rust-lang.org/cargo/appendix/glossary.html#target) as per your machine's architecture):
+
+```sh
+# from project root
+cargo run --release --features="native" --bin selfcaller-native --target aarch64-apple-darwin
+```
+
+This produces the `SystemTape` in both binary and debug formats.
+
+## Mozak-VM
+
+First, build the mozakvm binary:
+
+```sh
+# inside examples directory
+cargo build --release --bin selfcallerbin 
+```
+
+Test the ELF in mozak-vm using the below command. Note that you must run
+the native execution above to produce the system tape prior to running this.
+
+```sh
+# from project root
+MOZAK_STARK_DEBUG=true \
+    cargo run --bin mozak-cli -- prove-and-verify -vvv \
+    examples/target/riscv32im-mozak-mozakvm-elf/release/selfcallerbin \
+    --system-tape examples/selfcaller.tape.json \
+    --self-prog-id \
+    MZK-5b7b6135be198533f7c7ec46651216b762e6d47e69b408d1bc79d641f9ae06de;
+```

examples/selfcaller/core_logic.rs

@@ -0,0 +1,42 @@
+#![feature(restricted_std)]
+#![allow(unused_attributes)]
+extern crate alloc;
+
+use mozak_sdk::common::types::ProgramIdentifier;
+use rkyv::{Archive, Deserialize, Serialize};
+
+#[derive(Archive, Deserialize, Serialize, PartialEq, Eq, Clone)]
+#[cfg_attr(not(target_os = "mozakvm"), derive(Debug))]
+pub enum MethodArgs {
+    SelfCall(ProgramIdentifier, u8),
+}
+
+#[derive(Archive, Debug, Deserialize, Serialize, PartialEq, Eq, Clone)]
+pub enum MethodReturns {
+    SelfCall(u8),
+}
+
+// TODO: Remove later
+impl Default for MethodReturns {
+    fn default() -> Self { Self::SelfCall(0) }
+}
+
+#[allow(clippy::unit_arg)]
+pub fn dispatch(args: MethodArgs) -> MethodReturns {
+    match args {
+        MethodArgs::SelfCall(self_id, counter) => self_call(self_id, counter),
+    }
+}
+
+pub fn self_call(self_id: ProgramIdentifier, counter: u8) -> MethodReturns {
+    if counter == 0 {
+        return MethodReturns::SelfCall(0);
+    }
+    mozak_sdk::call_send(
+        self_id,
+        MethodArgs::SelfCall(self_id, counter - 1),
+        dispatch,
+    );
+
+    MethodReturns::SelfCall(counter)
+}

examples/selfcaller/main_mozak.rs

@@ -0,0 +1,17 @@
+#![no_main]
+#![allow(unused_attributes)]
+#![feature(restricted_std)]
+
+mod core_logic;
+
+use core_logic::{dispatch, MethodArgs, MethodReturns};
+use mozak_sdk::call_receive;
+
+pub fn main() {
+    while let Some((_caller, argument, return_)) = call_receive::<MethodArgs, MethodReturns>() {
+        assert!(dispatch(argument) == return_);
+    }
+}
+
+// We define `main()` to be the program's entry point.
+mozak_sdk::entry!(main);

examples/selfcaller/main_native.rs

@@ -0,0 +1,20 @@
+#![feature(restricted_std)]
+#![allow(unused_attributes)]
+
+mod core_logic;
+
+use mozak_sdk::common::types::ProgramIdentifier;
+
+use crate::core_logic::{dispatch, MethodArgs};
+
+fn main() {
+    let self_caller_program = ProgramIdentifier::new_from_rand_seed(5);
+
+    mozak_sdk::call_send(
+        self_caller_program,
+        MethodArgs::SelfCall(self_caller_program, 5),
+        dispatch,
+    );
+
+    mozak_sdk::native::dump_proving_files("selfcaller", self_caller_program);
+}

sdk/src/common/system.rs

@@ -11,6 +11,8 @@ use super::types::{
     CallTapeType, EventTapeType, PrivateInputTapeType, PublicInputTapeType, SystemTape,
 };
 #[cfg(target_os = "mozakvm")]
+use super::types::{SelfCallExtendedProgramIdentifier, SelfCallExtensionFlag};
+#[cfg(target_os = "mozakvm")]
 use crate::common::types::{CanonicalOrderedTemporalHints, CrossProgramCall, ProgramIdentifier};
 #[cfg(target_os = "mozakvm")]
 use crate::mozakvm::helpers::{
@@ -67,7 +69,10 @@ pub(crate) static mut SYSTEM_TAPE: Lazy<SystemTape> = Lazy::new(|| {
             private_input_tape: PrivateInputTapeType::default(),
             public_input_tape: PublicInputTapeType::default(),
             call_tape: CallTapeType {
-                self_prog_id: get_self_prog_id(),
+                extended_self_prog_id: SelfCallExtendedProgramIdentifier(
+                    get_self_prog_id(),
+                    SelfCallExtensionFlag::default(),
+                ),
                 cast_list: get_rkyv_deserialized!(Vec<ProgramIdentifier>, _mozak_cast_list),
                 reader: Some(get_rkyv_archived!(Vec<CrossProgramCall>, _mozak_call_tape)),
                 index: 0,

sdk/src/common/traits.rs

@@ -4,6 +4,7 @@ use rkyv::ser::{AllocSerializer, Composite};
 use rkyv::util::AlignedVec;
 use rkyv::{Archive, Deserialize, Serialize};
 
+use super::types::cross_program_call::SelfCallExtendedProgramIdentifier;
 use crate::common::types::{Event, ProgramIdentifier};
 
 pub trait RkyvSerializable = rkyv::Serialize<
@@ -14,9 +15,11 @@ pub trait CallReturn = ?Sized + Clone + Default + RkyvSerializable + Archive;
 
 /// A data struct that is aware of it's own ID
 pub trait SelfIdentify {
-    fn get_self_identity(&self) -> ProgramIdentifier;
+    fn get_self_identity(&self) -> SelfCallExtendedProgramIdentifier {
+        SelfCallExtendedProgramIdentifier::default()
+    }
     #[allow(dead_code)]
-    fn set_self_identity(&mut self, id: ProgramIdentifier);
+    fn set_self_identity(&mut self, _id: SelfCallExtendedProgramIdentifier) {}
 }
 
 /// `Call` trait provides methods `send` & `receive` to use an

sdk/src/common/types/cross_program_call.rs

@@ -1,3 +1,56 @@
+use super::ProgramIdentifier;
+
+#[derive(
+    Default,
+    Copy,
+    Clone,
+    Hash,
+    Eq,
+    PartialEq,
+    Ord,
+    PartialOrd,
+    rkyv::Archive,
+    rkyv::Serialize,
+    rkyv::Deserialize,
+)]
+#[cfg_attr(
+    not(target_os = "mozakvm"),
+    derive(Debug, serde::Serialize, serde::Deserialize)
+)]
+/// `SelfCallExtension` is a boolean value that differentiates the
+/// two program instances when a program calls itself under CPC. For
+/// example, if a program `P` in some function `f` wants to call another
+/// function `f'` in `P` under CPC regime, the caller and callee may form
+/// tuple as: `((P, 0), (P, 1))` or `((P, 1), (P, 0))`. The added digits
+/// help separate the two instances of call and hold no numeric significance
+/// apart from mere differentiators.
+pub struct SelfCallExtensionFlag(pub u8);
+
+impl SelfCallExtensionFlag {
+    /// Provides a flag different from given flag. In essence, if `0` is
+    /// provided, `1` is returned and vice versa
+    #[must_use]
+    pub fn differentiate_from(flag: Self) -> Self { Self(1 - flag.0) }
+}
+
+#[derive(
+    Default,
+    Clone,
+    Hash,
+    Eq,
+    PartialEq,
+    Ord,
+    PartialOrd,
+    rkyv::Archive,
+    rkyv::Serialize,
+    rkyv::Deserialize,
+)]
+#[cfg_attr(
+    not(target_os = "mozakvm"),
+    derive(Debug, serde::Serialize, serde::Deserialize)
+)]
+pub struct SelfCallExtendedProgramIdentifier(pub ProgramIdentifier, pub SelfCallExtensionFlag);
+
 #[derive(
     Default, Clone, Hash, PartialEq, PartialOrd, rkyv::Archive, rkyv::Serialize, rkyv::Deserialize,
 )]
@@ -7,8 +60,8 @@
 )]
 #[allow(clippy::pub_underscore_fields)]
 pub struct CrossProgramCall {
-    pub caller: super::ProgramIdentifier,
-    pub callee: super::ProgramIdentifier,
+    pub caller: SelfCallExtendedProgramIdentifier,
+    pub callee: SelfCallExtendedProgramIdentifier,
     pub argument: super::RawMessage,
     pub return_: super::RawMessage,
 }

sdk/src/common/types/mod.rs

@@ -7,7 +7,9 @@ pub(crate) mod state_address;
 pub(crate) mod state_object;
 pub(crate) mod system_tape;
 
-pub use cross_program_call::CrossProgramCall;
+pub use cross_program_call::{
+    CrossProgramCall, SelfCallExtendedProgramIdentifier, SelfCallExtensionFlag,
+};
 pub use event::{CanonicalEvent, CanonicalOrderedTemporalHints, Event, EventType};
 pub use poseidon2hash::Poseidon2Hash;
 pub use program_identifier::ProgramIdentifier;