ergotree-ir Add support for explicit type args, update Box.getReg sig…

…nature

Explicit type args are useful for many 6.0 methods where type can not be inferred from arguments, such as Box.getReg[T]

ergotree-ir/src/mir/method_call.rs

@@ -1,6 +1,9 @@
+use std::collections::HashMap;
+
 use crate::serialization::op_code::OpCode;
 use crate::types::smethod::SMethod;
 use crate::types::stype::SType;
+use crate::types::stype_param::STypeVar;
 
 use super::expr::Expr;
 use super::expr::InvalidArgumentError;
@@ -24,16 +27,36 @@ pub struct MethodCall {
     pub method: SMethod,
     /// Arguments passed to the method on invocation
     pub args: Vec<Expr>,
+    /// Arguments that cannot be inferred from function signature, such as Box.getReg[T]()
+    pub explicit_type_args: HashMap<STypeVar, SType>,
 }
 
 impl MethodCall {
-    /// Create new object, returns an error if any of the requirements failed
-    pub fn new(obj: Expr, method: SMethod, args: Vec<Expr>) -> Result<Self, InvalidArgumentError> {
+    fn new_inner(
+        method: SMethod,
+        args: Vec<Expr>,
+        obj: Expr,
+        explicit_type_args: HashMap<STypeVar, SType>,
+    ) -> Result<MethodCall, InvalidArgumentError> {
         if method.tpe().t_dom.len() != args.len() + 1 {
             return Err(InvalidArgumentError(format!(
                 "MethodCall: expected arguments count {} does not match provided arguments count {}",
                 method.tpe().t_dom.len(), args.len() + 1)));
         }
+        if method.method_raw.explicit_type_args.len() != explicit_type_args.len() {
+            return Err(InvalidArgumentError(format!("MethodCall: expected explicit type args count {} does not match provided type args count {}",
+                method.method_raw.explicit_type_args.len(), explicit_type_args.len())));
+        }
+        if let Some(missing_tpe) = method
+            .method_raw
+            .explicit_type_args
+            .iter()
+            .find(|tpe| !explicit_type_args.contains_key(&tpe))
+        {
+            return Err(InvalidArgumentError(format!(
+                "MethodCall: explicit_type_args does not include substitution for STypeVar {missing_tpe:?}",
+            )));
+        }
         let mut expected_types: Vec<SType> = vec![obj.tpe()];
         let arg_types: Vec<SType> = args.clone().into_iter().map(|a| a.tpe()).collect();
         expected_types.extend(arg_types);
@@ -53,9 +76,25 @@ impl MethodCall {
             obj: obj.into(),
             method,
             args,
+            explicit_type_args,
         })
     }
 
+    /// Create new object, returns an error if any of the requirements failed
+    pub fn new(obj: Expr, method: SMethod, args: Vec<Expr>) -> Result<Self, InvalidArgumentError> {
+        MethodCall::new_inner(method, args, obj, Default::default())
+    }
+
+    /// Create new object with explicit type args
+    pub fn with_type_args(
+        obj: Expr,
+        method: SMethod,
+        args: Vec<Expr>,
+        type_args: HashMap<STypeVar, SType>,
+    ) -> Result<Self, InvalidArgumentError> {
+        MethodCall::new_inner(method, args, obj, type_args)
+    }
+
     /// Type
     pub fn tpe(&self) -> SType {
         *self.method.tpe().t_range.clone()

ergotree-ir/src/pretty_printer/print.rs

@@ -518,6 +518,7 @@ impl Print for MethodCall {
                 obj: Box::new(obj),
                 method: self.method.clone(),
                 args,
+                explicit_type_args: self.explicit_type_args.clone(),
             },
         }
         .into())

ergotree-ir/src/serialization/method_call.rs

@@ -1,7 +1,11 @@
+use std::collections::HashMap;
+
 use crate::mir::expr::Expr;
 use crate::mir::method_call::MethodCall;
 use crate::types::smethod::MethodId;
 use crate::types::smethod::SMethod;
+use crate::types::stype::SType;
+use crate::types::stype_param::STypeVar;
 
 use super::sigma_byte_reader::SigmaByteRead;
 use super::sigma_byte_writer::SigmaByteWrite;
@@ -16,6 +20,11 @@ impl SigmaSerializable for MethodCall {
         self.method.method_id().sigma_serialize(w)?;
         self.obj.sigma_serialize(w)?;
         self.args.sigma_serialize(w)?;
+        for type_arg in &self.method.method_raw.explicit_type_args {
+            // Should not fail as existence of explicit type args is checked in constructor
+            let tpe = &self.explicit_type_args[type_arg];
+            tpe.sigma_serialize(w)?;
+        }
         Ok(())
     }
 
@@ -26,7 +35,20 @@ impl SigmaSerializable for MethodCall {
         let args = Vec::<Expr>::sigma_parse(r)?;
         let arg_types = args.iter().map(|arg| arg.tpe()).collect();
         let method = SMethod::from_ids(type_id, method_id)?.specialize_for(obj.tpe(), arg_types)?;
-        Ok(MethodCall::new(obj, method, args)?)
+        let explicit_type_args = method
+            .method_raw
+            .explicit_type_args
+            .iter()
+            .cloned()
+            .zip(std::iter::from_fn(|| Some(SType::sigma_parse(r))))
+            .map(|(tpe, res)| -> Result<(STypeVar, SType), SigmaParsingError> { Ok((tpe, res?)) })
+            .collect::<Result<HashMap<STypeVar, SType>, _>>()?;
+        Ok(MethodCall::with_type_args(
+            obj,
+            method.with_concrete_types(&explicit_type_args),
+            args,
+            explicit_type_args,
+        )?)
     }
 }
 

ergotree-ir/src/types/savltree.rs

@@ -78,6 +78,7 @@ lazy_static! {
             t_range: SType::SColl(Arc::new(SType::SByte)).into(),
             tpe_params: vec![],
         },
+        explicit_type_args: vec![]
     };
     /// AvlTree.digest
     pub static ref DIGEST_METHOD: SMethod =
@@ -93,6 +94,7 @@ lazy_static! {
             t_range: SType::SByte.into(),
             tpe_params: vec![],
         },
+        explicit_type_args: vec![]
     };
     /// AvlTree.enabledOperations
     pub static ref ENABLED_OPERATIONS_METHOD: SMethod =
@@ -108,6 +110,7 @@ lazy_static! {
             t_range: SType::SInt.into(),
             tpe_params: vec![],
         },
+        explicit_type_args: vec![]
     };
     /// AvlTree.keyLength
     pub static ref KEY_LENGTH_METHOD: SMethod =
@@ -123,6 +126,7 @@ lazy_static! {
             t_range: SType::SOption(Arc::new(SType::SInt)).into(),
             tpe_params: vec![],
         },
+        explicit_type_args: vec![]
     };
     /// AvlTree.valueLengthOpt
     pub static ref VALUE_LENGTH_OPT_METHOD: SMethod =
@@ -138,6 +142,7 @@ lazy_static! {
             t_range: SType::SBoolean.into(),
             tpe_params: vec![],
         },
+        explicit_type_args: vec![]
     };
     /// AvlTree.isInsertAllowed
     pub static ref IS_INSERT_ALLOWED_METHOD: SMethod =
@@ -153,6 +158,7 @@ lazy_static! {
             t_range: SType::SBoolean.into(),
             tpe_params: vec![],
         },
+        explicit_type_args: vec![]
     };
     /// AvlTree.isUpdateAllowed
     pub static ref IS_UPDATE_ALLOWED_METHOD: SMethod =
@@ -168,6 +174,7 @@ lazy_static! {
             t_range: SType::SBoolean.into(),
             tpe_params: vec![],
         },
+        explicit_type_args: vec![]
     };
     /// AvlTree.isRemoveAllowed
     pub static ref IS_REMOVE_ALLOWED_METHOD: SMethod =
@@ -183,6 +190,7 @@ lazy_static! {
             t_range: SType::SAvlTree.into(),
             tpe_params: vec![],
         },
+        explicit_type_args: vec![]
     };
     /// AvlTree.updateOperations
     pub static ref UPDATE_OPERATIONS_METHOD: SMethod =
@@ -201,6 +209,7 @@ lazy_static! {
             t_range: SType::SOption(SType::SColl(SType::SByte.into()).into()).into(),
             tpe_params: vec![],
         },
+        explicit_type_args: vec![]
     };
 
     /// AvlTree.get
@@ -220,6 +229,7 @@ lazy_static! {
             t_range: SType::SColl(SType::SOption(SType::SColl(SType::SByte.into()).into()).into()).into(),
             tpe_params: vec![],
         },
+        explicit_type_args: vec![]
     };
 
     /// AvlTree.getMany
@@ -248,6 +258,7 @@ lazy_static! {
             t_range: SType::SOption(Arc::new(SType::SAvlTree)).into(),
             tpe_params: vec![],
         },
+        explicit_type_args: vec![]
     };
     /// AvlTree.insert
     pub static ref INSERT_METHOD: SMethod =
@@ -270,6 +281,7 @@ lazy_static! {
             t_range: SType::SOption(Arc::new(SType::SAvlTree)).into(),
             tpe_params: vec![],
         },
+        explicit_type_args: vec![]
     };
     /// AvlTree.remove
     pub static ref REMOVE_METHOD: SMethod =
@@ -289,6 +301,7 @@ lazy_static! {
             t_range: SType::SBoolean.into(),
             tpe_params: vec![],
         },
+        explicit_type_args: vec![]
     };
     /// AvlTree.contains
     pub static ref CONTAINS_METHOD: SMethod =
@@ -316,6 +329,7 @@ lazy_static! {
             t_range: SType::SOption(Arc::new(SType::SAvlTree)).into(),
             tpe_params: vec![],
         },
+        explicit_type_args: vec![]
     };
     /// AvlTree.update
     pub static ref UPDATE_METHOD: SMethod =
@@ -331,6 +345,7 @@ lazy_static! {
             t_range: SType::SAvlTree.into(),
             tpe_params: vec![],
         },
+        explicit_type_args: vec![]
     };
     /// AvlTree.updateDigest
     pub static ref UPDATE_DIGEST_METHOD: SMethod =

ergotree-ir/src/types/sbox.rs

@@ -42,6 +42,7 @@ lazy_static! {
             t_range: Box::new(SType::SLong),
             tpe_params: vec![],
         },
+        explicit_type_args: vec![],
     };
     /// Box.value
     pub static ref VALUE_METHOD: SMethod = SMethod::new(STypeCompanion::Box, VALUE_METHOD_DESC.clone(),);
@@ -52,10 +53,11 @@ lazy_static! {
         method_id: GET_REG_METHOD_ID,
         name: "getReg",
         tpe: SFunc {
-            t_dom: vec![SType::SBox, SType::SByte],
+            t_dom: vec![SType::SBox, SType::SInt],
             t_range: SType::SOption(Arc::new(STypeVar::t().into())).into(),
             tpe_params: vec![],
         },
+        explicit_type_args: vec![STypeVar::t()]
     };
     /// Box.getReg
     pub static ref GET_REG_METHOD: SMethod =
@@ -75,6 +77,7 @@ lazy_static! {
                     ).into())).into(),
             tpe_params: vec![],
         },
+        explicit_type_args: vec![],
     };
     /// Box.tokens
     pub static ref TOKENS_METHOD: SMethod =
@@ -83,6 +86,11 @@ lazy_static! {
 
 #[cfg(test)]
 mod tests {
+    use crate::{
+        mir::{constant::Constant, global_vars::GlobalVars, method_call::MethodCall},
+        serialization::SigmaSerializable,
+    };
+
     use super::*;
 
     #[test]
@@ -91,4 +99,20 @@ mod tests {
         assert!(SMethod::from_ids(TYPE_CODE, GET_REG_METHOD_ID).map(|e| e.name()) == Ok("getReg"));
         assert!(SMethod::from_ids(TYPE_CODE, TOKENS_METHOD_ID).map(|e| e.name()) == Ok("tokens"));
     }
+
+    #[test]
+    fn test_getreg_serialization_roundtrip() {
+        let type_args = std::iter::once((STypeVar::t(), SType::SInt)).collect();
+        let mc = MethodCall::with_type_args(
+            GlobalVars::SelfBox.into(),
+            GET_REG_METHOD.clone().with_concrete_types(&type_args),
+            vec![Constant::from(4i32).into()],
+            type_args,
+        )
+        .unwrap();
+        assert_eq!(
+            MethodCall::sigma_parse_bytes(&mc.sigma_serialize_bytes().unwrap()).unwrap(),
+            mc
+        );
+    }
 }

ergotree-ir/src/types/scoll.rs

@@ -57,6 +57,7 @@ lazy_static! {
             t_range: SType::SInt.into(),
             tpe_params: vec![],
         },
+        explicit_type_args: vec![]
     };
     /// Coll.indexOf
     pub static ref INDEX_OF_METHOD: SMethod = SMethod::new(STypeCompanion::Coll, INDEX_OF_METHOD_DESC.clone());
@@ -76,6 +77,7 @@ lazy_static! {
                 ],
             SType::SColl(SType::STypeVar(STypeVar::ov()).into()),
         ),
+        explicit_type_args: vec![]
     };
     /// Coll.flatMap
     pub static ref FLATMAP_METHOD: SMethod = SMethod::new(STypeCompanion::Coll, FLATMAP_METHOD_DESC.clone());
@@ -93,7 +95,8 @@ lazy_static! {
             SType::SColl(SType::STuple(STuple::pair(
                 STypeVar::t().into(), STypeVar::iv().into()
             )).into())
-        )
+        ),
+        explicit_type_args: vec![]
     };
     /// Coll.zip
     pub static ref ZIP_METHOD: SMethod = SMethod::new(STypeCompanion::Coll, ZIP_METHOD_DESC.clone());
@@ -108,7 +111,8 @@ lazy_static! {
                 SType::SColl(SType::STypeVar(STypeVar::t()).into()),
             ],
             SType::SColl(SType::SInt.into())
-        )
+        ),
+        explicit_type_args: vec![]
     };
     /// Coll.indices
     pub static ref INDICES_METHOD: SMethod = SMethod::new(STypeCompanion::Coll, INDICES_METHOD_DESC.clone());
@@ -126,7 +130,8 @@ lazy_static! {
                 SType::SInt,
             ],
             SType::SColl(SType::STypeVar(STypeVar::t()).into())
-        )
+        ),
+        explicit_type_args: vec![]
     };
     /// Coll.patch
     pub static ref PATCH_METHOD: SMethod = SMethod::new(STypeCompanion::Coll, PATCH_METHOD_DESC.clone());
@@ -144,7 +149,8 @@ lazy_static! {
 
             ],
             SType::SColl(SType::STypeVar(STypeVar::t()).into())
-        )
+        ),
+        explicit_type_args: vec![]
     };
     /// Coll.updated
     pub static ref UPDATED_METHOD: SMethod = SMethod::new(STypeCompanion::Coll, UPDATED_METHOD_DESC.clone());
@@ -162,7 +168,8 @@ lazy_static! {
 
             ],
             SType::SColl(SType::STypeVar(STypeVar::t()).into())
-        )
+        ),
+        explicit_type_args: vec![]
     };
     /// Coll.updateMany
     pub static ref UPDATE_MANY_METHOD: SMethod = SMethod::new(STypeCompanion::Coll, UPDATE_MANY_METHOD_DESC.clone());

ergotree-ir/src/types/sglobal.rs

@@ -33,6 +33,7 @@ lazy_static! {
             t_range: SType::SGroupElement.into(),
             tpe_params: vec![],
         },
+        explicit_type_args: vec![]
     };
      /// GLOBAL.GroupGenerator
     pub static ref GROUP_GENERATOR_METHOD: SMethod = SMethod::new(STypeCompanion::Global, GROUP_GENERATOR_METHOD_DESC.clone(),);
@@ -52,6 +53,7 @@ lazy_static! {
             t_range: SType::SColl(SType::SByte.into()).into(),
             tpe_params: vec![],
         },
+        explicit_type_args: vec![]
     };
      /// GLOBAL.xor
     pub static ref XOR_METHOD: SMethod = SMethod::new(STypeCompanion::Global, XOR_METHOD_DESC.clone(),);