Panic if map's item length is not static

tasm-lib/src/data_type.rs

@@ -50,6 +50,24 @@ impl Display for StructType {
 }
 
 impl DataType {
+    /// See [`BFieldCodec::static_length`].
+    pub(crate) fn static_length(&self) -> Option<usize> {
+        match self {
+            DataType::Bool => bool::static_length(),
+            DataType::U32 => u32::static_length(),
+            DataType::U64 => u64::static_length(),
+            DataType::U128 => u128::static_length(),
+            DataType::Bfe => BFieldElement::static_length(),
+            DataType::Xfe => XFieldElement::static_length(),
+            DataType::Digest => Digest::static_length(),
+            DataType::List(_) => None,
+            DataType::Array(a) => Some(a.length * a.element_type.static_length()?),
+            DataType::Tuple(t) => t.iter().map(|dt| dt.static_length()).sum(),
+            DataType::VoidPointer => None,
+            DataType::StructRef(s) => s.fields.iter().map(|(_, dt)| dt.static_length()).sum(),
+        }
+    }
+
     /// Return a string which can be used as part of function labels in Triton-VM
     pub fn label_friendly_name(&self) -> String {
         match self {
@@ -92,8 +110,8 @@ impl DataType {
             DataType::Digest => Digest::LEN,
             DataType::List(_) => 1,
             DataType::Array(_) => 1,
-            DataType::VoidPointer => 1,
             DataType::Tuple(t) => t.iter().map(|dt| dt.stack_size()).sum(),
+            DataType::VoidPointer => 1,
             DataType::StructRef(_) => 1,
         }
     }
@@ -492,4 +510,100 @@ mod tests {
         let read = Literal::pop_from_stack(literal.data_type(), &mut vm_state.op_stack.stack);
         assert_eq!(literal, read);
     }
+
+    #[test]
+    fn static_lengths_match_up_for_vectors() {
+        assert_eq!(
+            <Vec<XFieldElement>>::static_length(),
+            DataType::List(Box::new(DataType::Xfe)).static_length()
+        );
+    }
+
+    #[test]
+    fn static_lengths_match_up_for_array_with_static_length_data_type() {
+        assert_eq!(
+            <[BFieldElement; 42]>::static_length(),
+            DataType::Array(Box::new(ArrayType {
+                element_type: DataType::Bfe,
+                length: 42
+            }))
+            .static_length()
+        );
+    }
+
+    #[test]
+    fn static_lengths_match_up_for_array_with_dynamic_length_data_type() {
+        assert_eq!(
+            <[Vec<BFieldElement>; 42]>::static_length(),
+            DataType::Array(Box::new(ArrayType {
+                element_type: DataType::List(Box::new(DataType::Bfe)),
+                length: 42
+            }))
+            .static_length()
+        );
+    }
+
+    #[test]
+    fn static_lengths_match_up_for_tuple_with_only_static_length_types() {
+        assert_eq!(
+            <(XFieldElement, BFieldElement)>::static_length(),
+            DataType::Tuple(vec![DataType::Xfe, DataType::Bfe]).static_length()
+        );
+    }
+
+    #[test]
+    fn static_lengths_match_up_for_tuple_with_dynamic_length_types() {
+        assert_eq!(
+            <(XFieldElement, Vec<BFieldElement>)>::static_length(),
+            DataType::Tuple(vec![DataType::Xfe, DataType::List(Box::new(DataType::Bfe))])
+                .static_length()
+        );
+    }
+
+    #[test]
+    fn static_length_of_void_pointer_is_unknown() {
+        assert!(DataType::VoidPointer.static_length().is_none());
+    }
+
+    #[test]
+    fn static_lengths_match_up_for_struct_with_only_static_length_types() {
+        #[derive(Debug, Clone, BFieldCodec)]
+        struct StructTyStatic {
+            u32: u32,
+            u64: u64,
+        }
+
+        let struct_ty_static = StructType {
+            name: "struct".to_owned(),
+            fields: vec![
+                ("u32".to_owned(), DataType::U32),
+                ("u64".to_owned(), DataType::U64),
+            ],
+        };
+        assert_eq!(
+            StructTyStatic::static_length(),
+            DataType::StructRef(struct_ty_static).static_length()
+        );
+    }
+
+    #[test]
+    fn static_lengths_match_up_for_struct_with_dynamic_length_types() {
+        #[derive(Debug, Clone, BFieldCodec)]
+        struct StructTyDyn {
+            digest: Digest,
+            list: Vec<BFieldElement>,
+        }
+
+        let struct_ty_dyn = StructType {
+            name: "struct".to_owned(),
+            fields: vec![
+                ("digest".to_owned(), DataType::Digest),
+                ("list".to_owned(), DataType::List(Box::new(DataType::Bfe))),
+            ],
+        };
+        assert_eq!(
+            StructTyDyn::static_length(),
+            DataType::StructRef(struct_ty_dyn).static_length()
+        );
+    }
 }

tasm-lib/src/list/higher_order/map.rs

@@ -74,12 +74,24 @@ impl<const NUM_INPUT_LISTS: usize> ChainMap<NUM_INPUT_LISTS> {
     ///
     /// - if the input type takes up [`OpStackElement::COUNT`] or more words
     /// - if the output type takes up [`OpStackElement::COUNT`]` - 1` or more words
+    /// - if the input type does not have a [static length][len]
+    /// - if the output type does not have a [static length][len]
+    ///
+    /// [len]: BFieldCodec::static_length
     pub fn new(f: InnerFunction) -> Self {
         // need instruction `place {input_type.stack_size()}`
-        assert!(f.domain().stack_size() < OpStackElement::COUNT);
+        let input_len = f
+            .domain()
+            .static_length()
+            .expect("input type's encoding length must be static");
+        assert!(input_len < OpStackElement::COUNT);
 
         // need instruction `pick {output_type.stack_size() + 1}`
-        assert!(f.range().stack_size() + 1 < OpStackElement::COUNT);
+        let output_len = f
+            .range()
+            .static_length()
+            .expect("output type's encoding length must be static");
+        assert!(output_len + 1 < OpStackElement::COUNT);
 
         Self { f }
     }
@@ -757,6 +769,17 @@ mod tests {
         test_case::<11>(guard);
         test_case::<12>(guard);
     }
+
+    #[test]
+    #[should_panic(expected = "static")]
+    fn mapping_over_dynamic_length_items_causes_a_panic() {
+        let raw_code = InnerFunction::RawCode(RawCode::new(
+            triton_asm!(irrelevant_label: return),
+            DataType::List(Box::new(DataType::Bfe)),
+            DataType::Digest,
+        ));
+        let _map = Map::new(raw_code);
+    }
 }
 
 #[cfg(test)]