From de476c801b5748d47323aed7d6b4b45269a2679f Mon Sep 17 00:00:00 2001
From: Kasey White <kwhitemsg@gmail.com>
Date: Sun, 3 Jul 2022 14:57:39 -0400
Subject: [PATCH] Add costing to steps

---
 crates/uplc/src/interpreter.rs | 106 ++++++++++++++++++++++++++-------
 1 file changed, 84 insertions(+), 22 deletions(-)

diff --git a/crates/uplc/src/interpreter.rs b/crates/uplc/src/interpreter.rs
index 3e109624..d9d97383 100644
--- a/crates/uplc/src/interpreter.rs
+++ b/crates/uplc/src/interpreter.rs
@@ -1,8 +1,12 @@
+use flat_rs::de;
+
 use crate::{
     ast::{Constant, NamedDeBruijn, Term},
     builtins::DefaultFunction,
 };
 
+#[repr(u8)]
+#[derive(Debug, Clone, PartialEq, Copy)]
 enum StepKind {
     BConst = 0,
     BVar = 1,
@@ -12,6 +16,25 @@ enum StepKind {
     BForce = 5,
     BBuiltin = 6,
 }
+impl TryFrom<u8> for StepKind {
+    type Error = de::Error;
+
+    fn try_from(value: u8) -> Result<Self, Self::Error> {
+        match value {
+            0 => Ok(StepKind::BConst),
+            1 => Ok(StepKind::BVar),
+            2 => Ok(StepKind::BLamAbs),
+            3 => Ok(StepKind::BApply),
+            4 => Ok(StepKind::BDelay),
+            5 => Ok(StepKind::BForce),
+            6 => Ok(StepKind::BBuiltin),
+            v => Err(de::Error::Message(format!(
+                "Default Function not found - {}",
+                v
+            ))),
+        }
+    }
+}
 
 enum ExBudgetCategory {
     BStep(StepKind),
@@ -25,6 +48,17 @@ struct ExBudget {
     cpu: i32,
 }
 
+impl std::ops::Mul<i32> for ExBudget {
+    type Output = ExBudget;
+
+    fn mul(self, rhs: i32) -> ExBudget {
+        ExBudget {
+            mem: self.mem * rhs,
+            cpu: self.cpu * rhs,
+        }
+    }
+}
+
 enum CekValue {
     VCon(Constant),
     VDelay(Term<NamedDeBruijn>, Vec<CekValue>),
@@ -47,44 +81,35 @@ enum Context {
 // For now let's just start with running cek on term with generic params
 fn run_cek_debruijn(term: Term<NamedDeBruijn>) -> (Term<NamedDeBruijn>, usize, Vec<String>) {
     //paramerterize this
-    let initial_budget = ExBudget {
+    let mut initial_budget = ExBudget {
         mem: 1000,
         cpu: 1000,
     };
-
-    let evaluation = enter_compute_cek(Context::NoFrame, Vec::new(), term);
+    let startup_budget = ExBudget { mem: 10, cpu: 10 };
+    spend_budget_cek(
+        &mut initial_budget,
+        ExBudgetCategory::BStartup,
+        startup_budget,
+    );
+    let evaluation = enter_compute_cek(Context::NoFrame, Vec::new(), term, &mut initial_budget);
     todo!()
 }
-
-fn spend_budget_cek(
-    current_budget: ExBudget,
-    category: ExBudgetCategory,
-    spend_budget: ExBudget,
-) -> ExBudget {
-    let new_budget = ExBudget {
-        mem: current_budget.mem - spend_budget.mem,
-        cpu: current_budget.cpu - spend_budget.cpu,
-    };
-    if new_budget.mem < 0 || new_budget.cpu < 0 {
-        panic!("Budget exhausted {:?}", new_budget);
-    }
-    new_budget
-}
-
 fn enter_compute_cek(
     frame: Context,
     env: Vec<CekValue>,
     term: Term<NamedDeBruijn>,
+    current_budget: &mut ExBudget,
 ) -> (Term<NamedDeBruijn>, usize, Vec<u32>) {
     //enter compute
-    compute_cek(Vec::with_capacity(8), frame, env, term)
+    compute_cek(&mut Vec::with_capacity(8), frame, env, term, current_budget)
 }
 
 fn compute_cek(
-    mut unbudgeted_steps: Vec<u32>,
+    unbudgeted_steps: &mut Vec<u32>,
     frame: Context,
     env: Vec<CekValue>,
     term: Term<NamedDeBruijn>,
+    current_budget: &mut ExBudget,
 ) -> (Term<NamedDeBruijn>, usize, Vec<u32>) {
     //TODO: parameterize slippage
     let slippage = 200;
@@ -92,7 +117,9 @@ fn compute_cek(
         a @ Term::Var(_) => {
             unbudgeted_steps[2] += 1;
             unbudgeted_steps[7] += 1;
-            //TODO spend budget one slippage is met
+            if unbudgeted_steps[7] >= slippage {
+                spend_accumulated_budget_cek(unbudgeted_steps, current_budget);
+            }
             todo!()
         }
         Term::Delay(_) => todo!(),
@@ -108,3 +135,38 @@ fn compute_cek(
     }
 }
 
+fn spend_accumulated_budget_cek(unbudgeted_steps: &mut Vec<u32>, current_budget: &mut ExBudget) {
+    //only spend each step kind and not total which is index
+    for i in 0..unbudgeted_steps.len() - 1 {
+        let step = StepKind::try_from(i as u8).unwrap();
+        spend_budget_cek(
+            current_budget,
+            ExBudgetCategory::BStep(step),
+            get_cost_by_step(step) * unbudgeted_steps[i] as i32,
+        );
+    }
+}
+
+fn spend_budget_cek(
+    current_budget: &mut ExBudget,
+    category: ExBudgetCategory,
+    spend_budget: ExBudget,
+) {
+    current_budget.mem -= spend_budget.mem;
+    current_budget.cpu -= spend_budget.cpu;
+    if current_budget.mem < 0 || current_budget.cpu < 0 {
+        panic!("Budget exhausted {:?}", current_budget);
+    }
+}
+
+fn get_cost_by_step(step: StepKind) -> ExBudget {
+    match (step) {
+        StepKind::BConst => ExBudget { mem: 100, cpu: 100 },
+        StepKind::BVar => ExBudget { mem: 100, cpu: 100 },
+        StepKind::BLamAbs => ExBudget { mem: 100, cpu: 100 },
+        StepKind::BApply => ExBudget { mem: 100, cpu: 100 },
+        StepKind::BDelay => ExBudget { mem: 100, cpu: 100 },
+        StepKind::BForce => ExBudget { mem: 100, cpu: 100 },
+        StepKind::BBuiltin => ExBudget { mem: 100, cpu: 100 },
+    }
+}