From cd441f439048b6a9371d172ca010a503a3bf18ae Mon Sep 17 00:00:00 2001
From: Yan Chen <yan.chen@dfinity.org>
Date: Sat, 9 Sep 2023 10:32:20 -0700
Subject: [PATCH] record and variant macro

---
 rust/candid/src/parser/typing.rs  |  7 ++++++-
 rust/candid/src/types/internal.rs | 32 +++++++++++++++++++++++++++----
 rust/candid/tests/parse_value.rs  | 11 +++++++++++
 3 files changed, 45 insertions(+), 5 deletions(-)

diff --git a/rust/candid/src/parser/typing.rs b/rust/candid/src/parser/typing.rs
index ec2d831d..1b126d59 100644
--- a/rust/candid/src/parser/typing.rs
+++ b/rust/candid/src/parser/typing.rs
@@ -248,7 +248,12 @@ pub fn check_prog(te: &mut TypeEnv, prog: &IDLProg) -> Result<Option<Type>> {
     check_actor(&env, &prog.actor)
 }
 /// Type check init args extracted from canister metadata candid:args.
-pub fn check_init_args(te: &mut TypeEnv, main_env: &TypeEnv, prog: &IDLInitArgs) -> Result<Vec<Type>> {
+/// Need to provide `main_env`, because init args may refer to variables from the main did file.
+pub fn check_init_args(
+    te: &mut TypeEnv,
+    main_env: &TypeEnv,
+    prog: &IDLInitArgs,
+) -> Result<Vec<Type>> {
     let mut env = Env { te, pre: false };
     check_decs(&mut env, &prog.decs)?;
     env.te.merge(main_env)?;
diff --git a/rust/candid/src/types/internal.rs b/rust/candid/src/types/internal.rs
index 9cfa5a16..b93f407a 100644
--- a/rust/candid/src/types/internal.rs
+++ b/rust/candid/src/types/internal.rs
@@ -356,9 +356,9 @@ impl fmt::Display for Field {
 /// Construct a field type, which can be used in `TypeInner::Record` and `TypeInner::Variant`.
 ///
 /// `field!{ a: TypeInner::Nat.into() }` expands to `Field { id: Label::Named("a"), ty: ... }`
-/// `field!{ 0: TypeInner::Nat.into() }` expands to `Field { id: Label::Id(0), ty: ... }`
+/// `field!{ 0: Nat::ty() }` expands to `Field { id: Label::Id(0), ty: ... }`
 macro_rules! field {
-    { $id:tt : $ty:expr } => {
+    { $id:tt : $ty:expr } => {{
         $crate::types::internal::Field {
             id: match stringify!($id).parse::<u32>() {
                 Ok(id) => $crate::types::Label::Id(id),
@@ -366,7 +366,31 @@ macro_rules! field {
             }.into(),
             ty: $ty
         }
-    };
+    }}
+}
+#[macro_export]
+/// Construct a record type, e.g., `record!{ label: Nat::ty(); 42: String::ty() }`.
+macro_rules! record {
+    { $($id:tt : $ty:expr);* $(;)? } => {{
+        let mut fs: Vec<$crate::types::internal::Field> = vec![ $($crate::field!{$id : $ty}),* ];
+        fs.sort_unstable_by_key(|f| f.id.get_id());
+        if let Err(e) = $crate::utils::check_unique(fs.iter().map(|f| &f.id)) {
+            panic!("{e}");
+        }
+        Into::<$crate::types::Type>::into($crate::types::TypeInner::Record(fs))
+    }}
+}
+#[macro_export]
+/// Construct a variant type, e.g., `variant!{ tag: <()>::ty() }`.
+macro_rules! variant {
+    { $($id:tt : $ty:expr);* $(;)? } => {{
+        let mut fs: Vec<$crate::types::internal::Field> = vec![ $($crate::field!{$id : $ty}),* ];
+        fs.sort_unstable_by_key(|f| f.id.get_id());
+        if let Err(e) = $crate::utils::check_unique(fs.iter().map(|f| &f.id)) {
+            panic!("{e}");
+        }
+        Into::<$crate::types::Type>::into($crate::types::TypeInner::Variant(fs))
+    }}
 }
 
 #[derive(Debug, Clone, PartialEq, Eq, Hash)]
@@ -431,7 +455,7 @@ macro_rules! func {
 /// `service!{ "f": func!((HttpRequest) -> ()) }` expands to `Type(Rc::new(TypeInner::Service(...)))`
 macro_rules! service {
     { $($meth:tt : $ty:expr);* $(;)? } => {{
-        let mut ms = vec![ $(($meth.to_string(), $ty)),* ];
+        let mut ms: Vec<(String, $crate::types::Type)> = vec![ $(($meth.to_string(), $ty)),* ];
         ms.sort_unstable_by(|a, b| a.0.as_str().partial_cmp(b.0.as_str()).unwrap());
         if let Err(e) = $crate::utils::check_unique(ms.iter().map(|m| &m.0)) {
             panic!("{e}");
diff --git a/rust/candid/tests/parse_value.rs b/rust/candid/tests/parse_value.rs
index 57710bf5..027401db 100644
--- a/rust/candid/tests/parse_value.rs
+++ b/rust/candid/tests/parse_value.rs
@@ -1,5 +1,6 @@
 use candid::types::value::{IDLArgs, IDLField, IDLValue, VariantValue};
 use candid::types::{Label, Type, TypeEnv, TypeInner};
+use candid::{record, variant, CandidType, Nat};
 
 fn parse_args(input: &str) -> IDLArgs {
     input.parse().unwrap()
@@ -136,6 +137,12 @@ fn parse_optional_record() {
     let mut args =
         parse_args("(opt record {}, record { 1=42;44=\"test\"; 2=false }, variant { 5=null })");
     let typ = parse_type("record { 1: nat; 44: text; 2: bool }");
+    assert_eq!(
+        typ,
+        record! { 1: Nat::ty(); 44: String::ty(); 2: bool::ty() }
+    );
+    assert_eq!(args.args[0].value_ty(), TypeInner::Opt(record! {}).into());
+    assert_eq!(args.args[2].value_ty(), variant! { 5: <()>::ty() });
     args.args[1] = args.args[1]
         .annotate_type(true, &TypeEnv::new(), &typ)
         .unwrap();
@@ -180,6 +187,10 @@ fn parse_nested_record() {
     let typ = parse_type(
         "record {label: nat; 0x2b:record { test:text; \"opt\":text }; long_label: opt null }",
     );
+    assert_eq!(
+        typ,
+        record! {label: Nat::ty(); 43: record!{ test: String::ty(); opt: String::ty() }; long_label: Option::<()>::ty(); }
+    );
     args.args[0] = args.args[0]
         .annotate_type(true, &TypeEnv::new(), &typ)
         .unwrap();