feat: most of flat encoding minus constant

crates/uplc/src/data.rs

@@ -2,7 +2,7 @@ use bumpalo::{collections::Vec as BumpVec, Bump};
 
 use crate::{
     constant::{integer_from, Constant, Integer},
-    flat::decode::Ctx,
+    flat::Ctx,
     machine::MachineError,
 };
 

crates/uplc/src/flat/data.rs

@@ -4,7 +4,7 @@ use rug::ops::NegAssign;
 
 use crate::data::PlutusData;
 
-use super::decode::Ctx;
+use super::Ctx;
 
 impl<'a, 'b> minicbor::decode::Decode<'b, Ctx<'a>> for &'a PlutusData<'a> {
     fn decode(

crates/uplc/src/flat/decode/mod.rs

@@ -4,3 +4,173 @@ mod error;
 pub use decoder::Ctx;
 pub use decoder::Decoder;
 pub use error::FlatDecodeError;
+
+use bumpalo::{
+    collections::{String as BumpString, Vec as BumpVec},
+    Bump,
+};
+
+use crate::{
+    constant::Constant,
+    program::{Program, Version},
+    term::Term,
+};
+
+use super::tag;
+use super::{
+    builtin,
+    tag::{BUILTIN_TAG_WIDTH, CONST_TAG_WIDTH, TERM_TAG_WIDTH},
+};
+
+pub fn decode<'a>(arena: &'a Bump, bytes: &[u8]) -> Result<&'a Program<'a>, FlatDecodeError> {
+    let mut decoder = Decoder::new(bytes);
+
+    let major = decoder.word()?;
+    let minor = decoder.word()?;
+    let patch = decoder.word()?;
+
+    let version = Version::new(arena, major, minor, patch);
+
+    let mut ctx = Ctx { arena };
+
+    let term = decode_term(&mut ctx, &mut decoder)?;
+
+    decoder.filler()?;
+
+    // TODO: probably should add a `finish()?` method that errors if bytes remain
+
+    Ok(Program::new(arena, version, term))
+}
+
+fn decode_term<'a>(
+    ctx: &mut Ctx<'a>,
+    decoder: &mut Decoder<'_>,
+) -> Result<&'a Term<'a>, FlatDecodeError> {
+    let tag = decoder.bits8(TERM_TAG_WIDTH)?;
+
+    match tag {
+        // Var
+        tag::VAR => Ok(Term::var(ctx.arena, decoder.word()?)),
+        // Delay
+        tag::DELAY => {
+            let term = decode_term(ctx, decoder)?;
+
+            Ok(term.delay(ctx.arena))
+        }
+        // Lambda
+        tag::LAMBDA => {
+            let term = decode_term(ctx, decoder)?;
+
+            Ok(term.lambda(ctx.arena, 0))
+        }
+        // Apply
+        tag::APPLY => {
+            let function = decode_term(ctx, decoder)?;
+            let argument = decode_term(ctx, decoder)?;
+
+            let term = function.apply(ctx.arena, argument);
+
+            Ok(term)
+        }
+        // Constant
+        tag::CONSTANT => {
+            let constant = decode_constant(ctx, decoder)?;
+
+            Ok(Term::constant(ctx.arena, constant))
+        }
+        // Force
+        tag::FORCE => {
+            let term = decode_term(ctx, decoder)?;
+
+            Ok(term.force(ctx.arena))
+        }
+        // Error
+        tag::ERROR => Ok(Term::error(ctx.arena)),
+        // Builtin
+        tag::BUILTIN => {
+            let builtin_tag = decoder.bits8(BUILTIN_TAG_WIDTH)?;
+
+            let function = builtin::try_from_tag(ctx.arena, builtin_tag)?;
+
+            let term = Term::builtin(ctx.arena, function);
+
+            Ok(term)
+        }
+        // Constr
+        tag::CONSTR => {
+            let tag = decoder.word()?;
+            let fields = decoder.list_with(ctx, decode_term)?;
+
+            let term = Term::constr(ctx.arena, tag, fields);
+
+            Ok(term)
+        }
+        // Case
+        tag::CASE => {
+            let constr = decode_term(ctx, decoder)?;
+            let branches = decoder.list_with(ctx, decode_term)?;
+
+            Ok(Term::case(ctx.arena, constr, branches))
+        }
+        _ => Err(FlatDecodeError::UnknownTermConstructor(tag)),
+    }
+}
+
+// BLS literals not supported
+fn decode_constant<'a>(
+    ctx: &mut Ctx<'a>,
+    d: &mut Decoder,
+) -> Result<&'a Constant<'a>, FlatDecodeError> {
+    let tags = decode_constant_tags(ctx, d)?;
+
+    match &tags.as_slice() {
+        [0] => {
+            let v = d.integer()?;
+
+            Ok(Constant::integer_from(ctx.arena, v))
+        }
+        [1] => {
+            let b = d.bytes(ctx.arena)?;
+
+            Ok(Constant::byte_string(ctx.arena, b))
+        }
+        [2] => {
+            let utf8_bytes = d.bytes(ctx.arena)?;
+
+            let s = BumpString::from_utf8(utf8_bytes)
+                .map_err(|e| FlatDecodeError::DecodeUtf8(e.utf8_error()))?;
+
+            Ok(Constant::string(ctx.arena, s))
+        }
+        [3] => Ok(Constant::unit(ctx.arena)),
+        [4] => {
+            let v = d.bit()?;
+
+            Ok(Constant::bool(ctx.arena, v))
+        }
+        [7, 5, rest @ ..] => todo!("list"),
+
+        [7, 7, 6, rest @ ..] => todo!("pair"),
+
+        [8] => {
+            let cbor = d.bytes(ctx.arena)?;
+
+            let data = minicbor::decode_with(&cbor, ctx)?;
+
+            Ok(Constant::data(ctx.arena, data))
+        }
+
+        x => Err(FlatDecodeError::UnknownConstantConstructor(x.to_vec())),
+    }
+}
+
+fn decode_constant_tags<'a>(
+    ctx: &mut Ctx<'a>,
+    d: &mut Decoder,
+) -> Result<BumpVec<'a, u8>, FlatDecodeError> {
+    d.list_with(ctx, |_arena, d| decode_constant_tag(d))
+}
+
+fn decode_constant_tag(d: &mut Decoder) -> Result<u8, FlatDecodeError> {
+    d.bits8(CONST_TAG_WIDTH)
+}

crates/uplc/src/flat/encode/encoder.rs

@@ -0,0 +1,149 @@
+use super::FlatEncodeError;
+
+#[derive(Default)]
+pub struct Encoder {
+    pub buffer: Vec<u8>,
+    // Int
+    used_bits: i64,
+    // Int
+    current_byte: u8,
+}
+
+impl Encoder {
+    /// Encode a unsigned integer of any size.
+    /// This is byte alignment agnostic.
+    /// We encode the 7 least significant bits of the unsigned byte. If the char
+    /// value is greater than 127 we encode a leading 1 followed by
+    /// repeating the above for the next 7 bits and so on.
+    pub fn word(&mut self, c: usize) -> &mut Self {
+        let mut d = c;
+        loop {
+            let mut w = (d & 127) as u8;
+            d >>= 7;
+
+            if d != 0 {
+                w |= 128;
+            }
+            self.bits(8, w);
+
+            if d == 0 {
+                break;
+            }
+        }
+
+        self
+    }
+
+    /// Encodes up to 8 bits of information and is byte alignment agnostic.
+    /// Uses unused bits in the current byte to write out the passed in byte
+    /// value. Overflows to the most significant digits of the next byte if
+    /// number of bits to use is greater than unused bits. Expects that
+    /// number of bits to use is greater than or equal to required bits by the
+    /// value. The param num_bits is i64 to match unused_bits type.
+    pub fn bits(&mut self, num_bits: i64, val: u8) -> &mut Self {
+        match (num_bits, val) {
+            (1, 0) => self.zero(),
+            (1, 1) => self.one(),
+            (2, 0) => {
+                self.zero();
+                self.zero();
+            }
+            (2, 1) => {
+                self.zero();
+                self.one();
+            }
+            (2, 2) => {
+                self.one();
+                self.zero();
+            }
+            (2, 3) => {
+                self.one();
+                self.one();
+            }
+            (_, _) => {
+                self.used_bits += num_bits;
+                let unused_bits = 8 - self.used_bits;
+                match unused_bits {
+                    0 => {
+                        self.current_byte |= val;
+                        self.next_word();
+                    }
+                    x if x > 0 => {
+                        self.current_byte |= val << x;
+                    }
+                    x => {
+                        let used = -x;
+                        self.current_byte |= val >> used;
+                        self.next_word();
+                        self.current_byte = val << (8 - used);
+                        self.used_bits = used;
+                    }
+                }
+            }
+        }
+
+        self
+    }
+
+    /// Encode a list of bytes with a function
+    /// 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 list_with<T>(
+        &mut self,
+        list: &[T],
+        encoder_func: for<'r> fn(&'r mut Encoder, &T) -> Result<(), FlatEncodeError>,
+    ) -> Result<&mut Self, FlatEncodeError> {
+        for item in list {
+            self.one();
+
+            encoder_func(self, item)?;
+        }
+
+        self.zero();
+
+        Ok(self)
+    }
+
+    /// A filler amount of end 0's followed by a 1 at the end of a byte.
+    /// Used to byte align the buffer by padding out the rest of the byte.
+    pub fn filler(&mut self) -> &mut Self {
+        self.current_byte |= 1;
+        self.next_word();
+
+        self
+    }
+
+    /// Write a 0 bit into the current byte.
+    /// Write out to buffer if last used bit in the current byte.
+    fn zero(&mut self) {
+        if self.used_bits == 7 {
+            self.next_word();
+        } else {
+            self.used_bits += 1;
+        }
+    }
+
+    /// Write a 1 bit into the current byte.
+    /// Write out to buffer if last used bit in the current byte.
+    fn one(&mut self) {
+        if self.used_bits == 7 {
+            self.current_byte |= 1;
+            self.next_word();
+        } else {
+            self.current_byte |= 128 >> self.used_bits;
+            self.used_bits += 1;
+        }
+    }
+
+    /// Write the current byte out to the buffer and begin next byte to write
+    /// out. Add current byte to the buffer and set current byte and used
+    /// bits to 0.
+    fn next_word(&mut self) {
+        self.buffer.push(self.current_byte);
+
+        self.current_byte = 0;
+        self.used_bits = 0;
+    }
+}

crates/uplc/src/flat/encode/error.rs

@@ -0,0 +1,7 @@
+use thiserror::Error;
+
+#[derive(Error, Debug)]
+pub enum FlatEncodeError {
+    #[error("Overflow detected, cannot fit {byte} in {num_bits} bits.")]
+    Overflow { byte: u8, num_bits: usize },
+}