flat encoding for list and pairs

Co-authored-by: rvcas <x@rvcas.dev>

add_integers.uplc

@@ -1,5 +1,5 @@
 (program
   1.0.0
-  [  (builtin subtractInteger ) (con integer 7) (con integer 4) ]
+  (con (list (list integer)) [[7], [5]])
 )
 

crates/flat/src/encode/encoder.rs

@@ -156,11 +156,14 @@ impl Encoder {
     /// This is byte alignment agnostic.
     /// If there are bytes in a list then write 1 bit followed by the functions encoding.
     /// After the last item write a 0 bit. If the list is empty only encode a 0 bit.
-    pub fn encode_list_with(
+    pub fn encode_list_with<T>(
         &mut self,
-        list: Vec<u8>,
+        list: &[T],
-        encoder_func: for<'r> fn(u8, &'r mut Encoder) -> Result<(), Error>,
+        encoder_func: for<'r> fn(&T, &'r mut Encoder) -> Result<(), Error>,
-    ) -> Result<&mut Self, Error> {
+    ) -> Result<&mut Self, Error>
+    where
+        T: Encode,
+    {
         for item in list {
             self.one();
             encoder_func(item, self)?;

crates/uplc/src/ast.rs

@@ -103,13 +103,42 @@ pub enum Constant {
     // tag: 4
     Bool(bool),
     // tag: 5
-    ProtoList(Vec<Constant>),
+    ProtoList(Type, Vec<Constant>),
     // tag: 6
-    ProtoPair(Box<Constant>, Box<Constant>),
+    ProtoPair(Type, Type, Box<Constant>, Box<Constant>),
+    // tag: 7
+    // Apply(Box<Constant>, Type),
     // tag: 8
     Data(Data),
 }
 
+#[derive(Debug, Clone, PartialEq)]
+pub enum Type {
+    Bool,
+    Integer,
+    String,
+    ByteString,
+    Unit,
+    List(Box<Type>),
+    Pair(Box<Type>, Box<Type>),
+    Data,
+}
+
+impl Display for Type {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        match self {
+            Type::Bool => write!(f, "bool"),
+            Type::Integer => write!(f, "integer"),
+            Type::String => write!(f, "string"),
+            Type::ByteString => write!(f, "bytestring"),
+            Type::Unit => write!(f, "unit"),
+            Type::List(t) => write!(f, "list {}", t),
+            Type::Pair(t1, t2) => write!(f, "pair {} {}", t1, t2),
+            Type::Data => write!(f, "data"),
+        }
+    }
+}
+
 #[derive(Debug, Clone, PartialEq)]
 pub enum Data {
     Constr(isize, Vec<Data>),

crates/uplc/src/flat.rs

@@ -7,7 +7,9 @@ use flat_rs::{
 };
 
 use crate::{
-    ast::{Constant, DeBruijn, FakeNamedDeBruijn, Name, NamedDeBruijn, Program, Term, Unique},
+    ast::{
+        Constant, DeBruijn, FakeNamedDeBruijn, Name, NamedDeBruijn, Program, Term, Type, Unique,
+    },
     builtins::DefaultFunction,
 };
 
@@ -161,7 +163,7 @@ where
     }
 }
 
-impl Encode for &Constant {
+impl Encode for Constant {
     fn encode(&self, e: &mut Encoder) -> Result<(), en::Error> {
         match self {
             // Integers are typically smaller so we save space
@@ -171,29 +173,93 @@ impl Encode for &Constant {
             // i.e. A 17 or greater length byte array loses efficiency being encoded as
             // a unsigned integer instead of a byte array
             Constant::Integer(i) => {
-                encode_constant(0, e)?;
+                encode_constant(&[0], e)?;
                 i.encode(e)?;
             }
 
             Constant::ByteString(bytes) => {
-                encode_constant(1, e)?;
+                encode_constant(&[1], e)?;
                 bytes.encode(e)?;
             }
             Constant::String(s) => {
-                encode_constant(2, e)?;
+                encode_constant(&[2], e)?;
                 s.encode(e)?;
             }
-            Constant::Unit => encode_constant(3, e)?,
+            Constant::Unit => encode_constant(&[3], e)?,
             Constant::Bool(b) => {
-                encode_constant(4, e)?;
+                encode_constant(&[4], e)?;
                 b.encode(e)?;
             }
+            Constant::ProtoList(typ, list) => {
+                let mut type_encode = vec![7, 5];
+
+                encode_type(typ, &mut type_encode);
+
+                encode_constant(&type_encode, e)?;
+
+                e.encode_list_with(list, encode_constant_value)?;
+            }
+            Constant::ProtoPair(type1, type2, a, b) => {
+                let mut type_encode = vec![7, 7, 6];
+
+                encode_type(type1, &mut type_encode);
+
+                encode_type(type2, &mut type_encode);
+
+                encode_constant(&type_encode, e)?;
+                encode_constant_value(a, e)?;
+                encode_constant_value(b, e)?;
+            }
+            Constant::Data(_) => {
+                encode_constant(&[8], e)?;
+                todo!()
+            }
         }
 
         Ok(())
     }
 }
 
+fn encode_constant_value(x: &Constant, e: &mut Encoder) -> Result<(), en::Error> {
+    match x {
+        Constant::Integer(x) => x.encode(e),
+        Constant::ByteString(b) => b.encode(e),
+        Constant::String(s) => s.encode(e),
+        Constant::Unit => Ok(()),
+        Constant::Bool(b) => b.encode(e),
+        Constant::ProtoList(_, list) => {
+            e.encode_list_with(list, encode_constant_value)?;
+            Ok(())
+        }
+        Constant::ProtoPair(_, _, a, b) => {
+            encode_constant_value(a, e)?;
+            encode_constant_value(b, e)?;
+            Ok(())
+        }
+        Constant::Data(_) => todo!(),
+    }
+}
+
+fn encode_type(typ: &Type, bytes: &mut Vec<u8>) {
+    match typ {
+        Type::Bool => bytes.push(4),
+        Type::Integer => bytes.push(0),
+        Type::String => bytes.push(2),
+        Type::ByteString => bytes.push(1),
+        Type::Unit => bytes.push(3),
+        Type::List(sub_typ) => {
+            bytes.extend(vec![7, 5]);
+            encode_type(sub_typ, bytes);
+        }
+        Type::Pair(type1, type2) => {
+            bytes.extend(vec![7, 7, 6]);
+            encode_type(type1, bytes);
+            encode_type(type2, bytes);
+        }
+        Type::Data => todo!(),
+    }
+}
+
 impl<'b> Decode<'b> for Constant {
     fn decode(d: &mut Decoder) -> Result<Self, de::Error> {
         match decode_constant(d)? {
@@ -392,8 +458,8 @@ fn safe_encode_bits(num_bits: u32, byte: u8, e: &mut Encoder) -> Result<(), en::
     }
 }
 
-pub fn encode_constant(tag: u8, e: &mut Encoder) -> Result<(), en::Error> {
+pub fn encode_constant(tag: &[u8], e: &mut Encoder) -> Result<(), en::Error> {
-    e.encode_list_with([tag].to_vec(), encode_constant_tag)?;
+    e.encode_list_with(tag, encode_constant_tag)?;
 
     Ok(())
 }
@@ -410,8 +476,8 @@ pub fn decode_constant(d: &mut Decoder) -> Result<u8, de::Error> {
     }
 }
 
-pub fn encode_constant_tag(tag: u8, e: &mut Encoder) -> Result<(), en::Error> {
+pub fn encode_constant_tag(tag: &u8, e: &mut Encoder) -> Result<(), en::Error> {
-    safe_encode_bits(CONST_TAG_WIDTH, tag, e)
+    safe_encode_bits(CONST_TAG_WIDTH, *tag, e)
 }
 
 pub fn decode_constant_tag(d: &mut Decoder) -> Result<u8, de::Error> {
@@ -422,7 +488,7 @@ pub fn decode_constant_tag(d: &mut Decoder) -> Result<u8, de::Error> {
 mod test {
     use flat_rs::Flat;
 
-    use crate::ast::Name;
+    use crate::ast::{Name, Type};
 
     use super::{Constant, Program, Term};
 
@@ -441,6 +507,58 @@ mod test {
         )
     }
 
+    #[test]
+    fn flat_encode_list_list_integer() {
+        let program = Program::<Name> {
+            version: (1, 0, 0),
+            term: Term::Constant(Constant::ProtoList(
+                Type::List(Box::new(Type::Integer)),
+                vec![
+                    Constant::ProtoList(Type::Integer, vec![Constant::Integer(7)]),
+                    Constant::ProtoList(Type::Integer, vec![Constant::Integer(5)]),
+                ],
+            )),
+        };
+
+        let bytes = program.to_flat().unwrap();
+
+        assert_eq!(
+            bytes,
+            vec![
+                0b00000001, 0b00000000, 0b00000000, 0b01001011, 0b11010110, 0b11110101, 0b10000011,
+                0b00001110, 0b01100001, 0b01000001
+            ]
+        )
+    }
+
+    #[test]
+    fn flat_encode_pair_pair_integer_bool_integer() {
+        let program = Program::<Name> {
+            version: (1, 0, 0),
+            term: Term::Constant(Constant::ProtoPair(
+                Type::Pair(Box::new(Type::Integer), Box::new(Type::Bool)),
+                Type::Integer,
+                Box::new(Constant::ProtoPair(
+                    Type::Integer,
+                    Type::Bool,
+                    Box::new(Constant::Integer(11)),
+                    Box::new(Constant::Bool(true)),
+                )),
+                Box::new(Constant::Integer(11)),
+            )),
+        };
+
+        let bytes = program.to_flat().unwrap();
+
+        assert_eq!(
+            bytes,
+            vec![
+                0b00000001, 0b00000000, 0b00000000, 0b01001011, 0b11011110, 0b11010111, 0b10111101,
+                0b10100001, 0b01001000, 0b00000101, 0b10100010, 0b11000001
+            ]
+        )
+    }
+
     #[test]
     fn flat_decode_integer() {
         let flat_encoded = vec![

crates/uplc/src/machine.rs

@@ -1,5 +1,5 @@
 use crate::{
-    ast::{Constant, NamedDeBruijn, Term},
+    ast::{Constant, NamedDeBruijn, Term, Type},
     builtins::DefaultFunction,
 };
 
@@ -10,7 +10,7 @@ mod runtime;
 use cost_model::{ExBudget, StepKind};
 pub use error::Error;
 
-use self::{cost_model::CostModel, error::Type, runtime::BuiltinRuntime};
+use self::{cost_model::CostModel, runtime::BuiltinRuntime};
 
 enum MachineStep {
     Return(Context, Value),
@@ -434,6 +434,9 @@ impl Value {
                 Constant::String(s) => s.chars().count() as i64,
                 Constant::Unit => 1,
                 Constant::Bool(_) => 1,
+                Constant::ProtoList(_, _) => todo!(),
+                Constant::ProtoPair(_, _, _, _) => todo!(),
+                Constant::Data(_) => todo!(),
             },
             Value::Delay(_, _) => 1,
             Value::Lambda { .. } => 1,
@@ -465,6 +468,11 @@ impl From<&Constant> for Type {
             Constant::String(_) => Type::String,
             Constant::Unit => Type::Unit,
             Constant::Bool(_) => Type::Bool,
+            Constant::ProtoList(t, _) => Type::List(Box::new(t.clone())),
+            Constant::ProtoPair(t1, t2, _, _) => {
+                Type::Pair(Box::new(t1.clone()), Box::new(t2.clone()))
+            }
+            Constant::Data(_) => todo!(),
         }
     }
 }

crates/uplc/src/machine/error.rs

@@ -1,8 +1,6 @@
-use std::fmt::Display;
-
 use thiserror::Error;
 
-use crate::ast::{NamedDeBruijn, Term};
+use crate::ast::{NamedDeBruijn, Term, Type};
 
 use super::{ExBudget, Value};
 
@@ -31,24 +29,3 @@ pub enum Error {
     #[error("The evaluation never reached a final state")]
     MachineNeverReachedDone,
 }
-
-#[derive(Debug, Clone, PartialEq)]
-pub enum Type {
-    Bool,
-    Integer,
-    String,
-    ByteString,
-    Unit,
-}
-
-impl Display for Type {
-    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
-        match self {
-            Type::Bool => write!(f, "bool"),
-            Type::Integer => write!(f, "integer"),
-            Type::String => write!(f, "string"),
-            Type::ByteString => write!(f, "bytestring"),
-            Type::Unit => write!(f, "unit"),
-        }
-    }
-}

crates/uplc/src/machine/runtime.rs

@@ -1,8 +1,10 @@
-use crate::{ast::Constant, builtins::DefaultFunction};
+use crate::{
+    ast::{Constant, Type},
+    builtins::DefaultFunction,
+};
 
 use super::{
     cost_model::{BuiltinCosts, ExBudget},
-    error::Type,
     Error, Value,
 };
 

crates/uplc/src/pretty.rs

@@ -170,6 +170,9 @@ impl Constant {
             Constant::Bool(b) => RcDoc::text("bool")
                 .append(RcDoc::line())
                 .append(RcDoc::text(if *b { "True" } else { "False" })),
+            Constant::ProtoList(_, _) => todo!(),
+            Constant::ProtoPair(_, _, _, _) => todo!(),
+            Constant::Data(_) => todo!(),
         }
     }
 }