From fa2aa0a3e8a0baec4bc2d288dedbf8d406975f44 Mon Sep 17 00:00:00 2001
From: microproofs <kwhitemsg@gmail.com>
Date: Fri, 20 Sep 2024 12:31:36 -0400
Subject: [PATCH] When we find bls constants in UPLC convert them to their a
 hoisted compressed form with an uncompress call

---
 crates/uplc/src/optimize.rs          |  1 +
 crates/uplc/src/optimize/shrinker.rs | 67 +++++++++++++++++++++++++++-
 2 files changed, 67 insertions(+), 1 deletion(-)

diff --git a/crates/uplc/src/optimize.rs b/crates/uplc/src/optimize.rs
index 550a020f..ef60c20a 100644
--- a/crates/uplc/src/optimize.rs
+++ b/crates/uplc/src/optimize.rs
@@ -6,6 +6,7 @@ pub mod shrinker;
 pub fn aiken_optimize_and_intern(program: Program<Name>) -> Program<Name> {
     program
         .inline_constr_ops()
+        .bls381_compressor()
         .builtin_force_reducer()
         .lambda_reducer()
         .inline_reducer()
diff --git a/crates/uplc/src/optimize/shrinker.rs b/crates/uplc/src/optimize/shrinker.rs
index 7913ad1e..72dbf68c 100644
--- a/crates/uplc/src/optimize/shrinker.rs
+++ b/crates/uplc/src/optimize/shrinker.rs
@@ -3,7 +3,7 @@ use crate::{
     ast::{Constant, Data, Name, NamedDeBruijn, Program, Term, Type},
     builder::{CONSTR_FIELDS_EXPOSER, CONSTR_INDEX_EXPOSER},
     builtins::DefaultFunction,
-    machine::cost_model::ExBudget,
+    machine::{cost_model::ExBudget, runtime::Compressable},
 };
 use indexmap::IndexMap;
 use itertools::Itertools;
@@ -1138,6 +1138,71 @@ impl Program<Name> {
         Program::<Name>::try_from(program).unwrap()
     }
 
+    pub fn bls381_compressor(self) -> Self {
+        let mut blst_p1_list = vec![];
+        let mut blst_p2_list = vec![];
+
+        let program = self.traverse_uplc_with(true, &mut |_id, term, _arg_stack, _scope| {
+            if let Term::Constant(con) = term {
+                match con.as_ref() {
+                    Constant::Bls12_381G1Element(blst_p1) => {
+                        if let Some(index) = blst_p1_list
+                            .iter()
+                            .position(|item| item == blst_p1.as_ref())
+                        {
+                            *term = Term::var(format!("blst_p1_index_{}", index));
+                        } else {
+                            blst_p1_list.push(blst_p1.as_ref().clone());
+                            *term = Term::var(format!("blst_p1_index_{}", blst_p1_list.len() - 1));
+                        }
+                    }
+                    Constant::Bls12_381G2Element(blst_p2) => {
+                        if let Some(index) = blst_p2_list
+                            .iter()
+                            .position(|item| item == blst_p2.as_ref())
+                        {
+                            *term = Term::var(format!("blst_p2_index_{}", index));
+                        } else {
+                            blst_p2_list.push(blst_p2.as_ref().clone());
+                            *term = Term::var(format!("blst_p2_index_{}", blst_p2_list.len() - 1));
+                        }
+                    }
+                    _ => (),
+                }
+            }
+        });
+        let mut term = program.term;
+
+        for (index, blst_p1) in blst_p1_list.into_iter().enumerate() {
+            let compressed = blst_p1.compress();
+
+            term = term
+                .lambda(format!("blst_p1_index_{}", index))
+                .apply(Term::bls12_381_g1_uncompress().apply(Term::byte_string(compressed)));
+        }
+
+        for (index, blst_p2) in blst_p2_list.into_iter().enumerate() {
+            let compressed = blst_p2.compress();
+
+            term = term
+                .lambda(format!("blst_p2_index_{}", index))
+                .apply(Term::bls12_381_g2_uncompress().apply(Term::byte_string(compressed)));
+        }
+
+        let mut program = Program {
+            version: program.version,
+            term,
+        };
+
+        let mut interner = CodeGenInterner::new();
+
+        interner.program(&mut program);
+
+        let program = Program::<NamedDeBruijn>::try_from(program).unwrap();
+
+        Program::<Name>::try_from(program).unwrap()
+    }
+
     pub fn identity_reducer(self) -> Self {
         let mut identity_applied_ids = vec![];
         self.traverse_uplc_with(true, &mut |id, term, mut arg_stack, _scope| {