feat: costing for fst and snd pair builtins

pragma-org · Oct 30, 2024 · e077864 · e077864
1 parent b2f0150
commit e077864
Show file tree

Hide file tree

Showing 5 changed files with 43 additions and 3 deletions.
diff --git a/crates/uplc/src/machine/cost_model/builtin_costs.rs b/crates/uplc/src/machine/cost_model/builtin_costs.rs
@@ -306,6 +306,14 @@ impl BuiltinCosts {
         )
     }
 
+    pub fn fst_pair(&self, args: [i64; 1]) -> ExBudget {
+        ExBudget::new(self.fst_pair.mem.cost(args), self.fst_pair.cpu.cost(args))
+    }
+
+    pub fn snd_pair(&self, args: [i64; 1]) -> ExBudget {
+        ExBudget::new(self.snd_pair.mem.cost(args), self.snd_pair.cpu.cost(args))
+    }
+
     pub fn v3() -> Self {
         Self {
             add_integer: TwoArgumentsCosting::new(

diff --git a/crates/uplc/src/machine/cost_model/value.rs b/crates/uplc/src/machine/cost_model/value.rs
@@ -37,6 +37,10 @@ pub fn string_ex_mem(s: &str) -> i64 {
     s.chars().count() as i64
 }
 
+pub fn pair_ex_mem(l: &Constant, r: &Constant) -> i64 {
+    constant_ex_mem(l) + constant_ex_mem(r)
+}
+
 pub fn value_ex_mem(v: &Value) -> i64 {
     match v {
         Value::Con(c) => constant_ex_mem(c),
@@ -57,7 +61,7 @@ pub fn constant_ex_mem(c: &Constant) -> i64 {
         Constant::ProtoList(_, items) => items
             .iter()
             .fold(0, |acc, constant| acc + constant_ex_mem(constant)),
-        Constant::ProtoPair(_, _, l, r) => constant_ex_mem(l) + constant_ex_mem(r),
+        Constant::ProtoPair(_, _, l, r) => pair_ex_mem(l, r),
         Constant::Data(d) => data_ex_mem(d),
         Constant::Bls12_381G1Element(_) => size_of::<blst::blst_p1>() as i64 / 8,
         Constant::Bls12_381G2Element(_) => size_of::<blst::blst_p2>() as i64 / 8,

diff --git a/crates/uplc/src/machine/runtime.rs b/crates/uplc/src/machine/runtime.rs
@@ -806,14 +806,28 @@ impl<'a> Machine<'a> {
             }
             DefaultFunction::Trace => todo!(),
             DefaultFunction::FstPair => {
-                let (_, _, first, _) = runtime.args[0].unwrap_pair()?;
+                let (_, _, first, second) = runtime.args[0].unwrap_pair()?;
+
+                let budget = self
+                    .costs
+                    .builtin_costs
+                    .fst_pair([cost_model::pair_ex_mem(first, second)]);
+
+                self.spend_budget(budget)?;
 
                 let value = Value::con(self.arena, first);
 
                 Ok(value)
             }
             DefaultFunction::SndPair => {
-                let (_, _, _, second) = runtime.args[0].unwrap_pair()?;
+                let (_, _, first, second) = runtime.args[0].unwrap_pair()?;
+
+                let budget = self
+                    .costs
+                    .builtin_costs
+                    .snd_pair([cost_model::pair_ex_mem(first, second)]);
+
+                self.spend_budget(budget)?;
 
                 let value = Value::con(self.arena, second);
 

diff --git a/crates/uplc/src/syn/term.rs b/crates/uplc/src/syn/term.rs
@@ -153,6 +153,8 @@ pub fn builtin_from_str<'a>(arena: &'a Bump, name: &str) -> Option<&'a Term<'a>>
         "encodeUtf8" => Some(Term::encode_utf8(arena)),
         "decodeUtf8" => Some(Term::decode_utf8(arena)),
         "chooseUnit" => Some(Term::choose_unit(arena)),
+        "fstPair" => Some(Term::fst_pair(arena)),
+        "sndPair" => Some(Term::snd_pair(arena)),
         _ => None,
     }
 }
diff --git a/crates/uplc/src/term.rs b/crates/uplc/src/term.rs
@@ -345,4 +345,16 @@ impl<'a> Term<'a> {
 
         Term::builtin(arena, fun)
     }
+
+    pub fn fst_pair(arena: &'a Bump) -> &'a Term<'a> {
+        let fun = arena.alloc(DefaultFunction::FstPair);
+
+        Term::builtin(arena, fun)
+    }
+
+    pub fn snd_pair(arena: &'a Bump) -> &'a Term<'a> {
+        let fun = arena.alloc(DefaultFunction::SndPair);
+
+        Term::builtin(arena, fun)
+    }
 }