From f04832a8bf46df822d76ecc279d4c31361dec03c Mon Sep 17 00:00:00 2001
From: Gaston Zanitti <gzanitti@gmail.com>
Date: Tue, 13 Aug 2024 17:59:58 -0300
Subject: [PATCH] WIP

---
 ast/src/analyzed/mod.rs              |  2 +-
 pil-analyzer/src/evaluator.rs        | 41 ++++++++++------------------
 pil-analyzer/src/traits_processor.rs | 11 ++++----
 3 files changed, 20 insertions(+), 34 deletions(-)

diff --git a/ast/src/analyzed/mod.rs b/ast/src/analyzed/mod.rs
index 437a82ada..5e2614a54 100644
--- a/ast/src/analyzed/mod.rs
+++ b/ast/src/analyzed/mod.rs
@@ -1326,7 +1326,7 @@ pub struct PolynomialReference {
     /// Guaranteed to be Some(_) after type checking is completed.
     pub type_args: Option<Vec<Type>>,
     ///
-    pub resolved_impls: BTreeMap<Vec<Type>, Box<Expression>>,
+    pub resolved_impls: BTreeMap<Vec<Type>, usize>,
 }
 
 #[derive(
diff --git a/pil-analyzer/src/evaluator.rs b/pil-analyzer/src/evaluator.rs
index e4d66a9a4..1e5a88d70 100644
--- a/pil-analyzer/src/evaluator.rs
+++ b/pil-analyzer/src/evaluator.rs
@@ -138,7 +138,6 @@ pub enum Value<'a, T> {
     Enum(&'a str, Option<Vec<Arc<Self>>>),
     BuiltinFunction(BuiltinFunction),
     Expression(AlgebraicExpression<T>),
-    TraitFunction(Closure<'a, T>),
 }
 
 impl<'a, T: FieldElement> From<T> for Value<'a, T> {
@@ -216,12 +215,6 @@ impl<'a, T: FieldElement> Value<'a, T> {
             Value::Enum(name, _) => name.to_string(),
             Value::BuiltinFunction(b) => format!("builtin_{b:?}"),
             Value::Expression(_) => "expr".to_string(),
-            Value::TraitFunction(trait_function) => {
-                format!(
-                    "trait_function<{}>",
-                    trait_function.type_args.values().format(", ")
-                )
-            }
         }
     }
 
@@ -377,11 +370,6 @@ impl<'a, T: Display> Display for Value<'a, T> {
             }
             Value::BuiltinFunction(b) => write!(f, "{b:?}"),
             Value::Expression(e) => write!(f, "{e}"),
-            Value::TraitFunction(trait_function) => write!(
-                f,
-                "trait_function<{}>",
-                trait_function.type_args.values().format(", ")
-            ),
         }
     }
 }
@@ -833,7 +821,7 @@ impl<'a, 'b, T: FieldElement, S: SymbolLookup<'a, T>> Evaluator<'a, 'b, T, S> {
                         .and_then(|type_args| poly.resolved_impls.get(type_args).as_ref().copied());
 
                     match impl_pos {
-                        Some(expr) => {
+                        Some(index) => {
                             let local_type_args = poly
                                 .type_args
                                 .clone()
@@ -849,16 +837,20 @@ impl<'a, 'b, T: FieldElement, S: SymbolLookup<'a, T>> Evaluator<'a, 'b, T, S> {
                                 })
                                 .unwrap();
 
-                            let Expression::LambdaExpression(_, body) = &expr.as_ref() else {
-                                unreachable!()
-                            };
+                            // let Expression::LambdaExpression(_, body) = &expr.as_ref() else {
+                            //     unreachable!()
+                            // };
 
-                            Value::Closure(Closure {
-                                lambda: body,
-                                environment: vec![],
-                                type_args: local_type_args.clone(),
-                            })
-                            .into()
+                            let symbol = self.symbols.lookup(&poly.name, &type_args)?;
+                            println!("{:?}", symbol);
+                            symbol
+
+                            // Value::Closure(Closure {
+                            //     lambda: body,
+                            //     environment: vec![],
+                            //     type_args: local_type_args.clone(),
+                            // })
+                            // .into()
                         }
                         None => self.symbols.lookup(&poly.name, &type_args)?,
                     }
@@ -1004,11 +996,6 @@ impl<'a, 'b, T: FieldElement, S: SymbolLookup<'a, T>> Evaluator<'a, 'b, T, S> {
                 lambda,
                 environment,
                 type_args,
-            })
-            | Value::TraitFunction(Closure {
-                lambda,
-                environment,
-                type_args,
             }) => {
                 assert_eq!(lambda.params.len(), arguments.len());
                 let matched_arguments =
diff --git a/pil-analyzer/src/traits_processor.rs b/pil-analyzer/src/traits_processor.rs
index f629dcea8..a25ea0d8a 100644
--- a/pil-analyzer/src/traits_processor.rs
+++ b/pil-analyzer/src/traits_processor.rs
@@ -17,7 +17,7 @@ impl<'a> TraitsProcessor<'a> {
 
     pub fn traits_resolution(
         &self,
-        references: &mut Vec<&mut Reference>,
+        references: &mut [&mut Reference],
         traits_functions_defs: HashMap<String, FunctionValueDefinition>,
     ) {
         for r in references.iter_mut() {
@@ -44,9 +44,8 @@ impl<'a> TraitsProcessor<'a> {
         &self,
         collected_ref: (String, &Vec<Type>),
         traits_functions_definition: &HashMap<String, FunctionValueDefinition>,
-    ) -> Option<(String, BTreeMap<Vec<Type>, Box<Expression>>)> {
+    ) -> Option<(String, BTreeMap<Vec<Type>, usize>)> {
         let mut resolved_impl_pos = BTreeMap::new();
-
         let Some(FunctionValueDefinition::TraitFunction(ref trait_decl, ref trait_fn)) =
             traits_functions_definition.get(&collected_ref.0)
         else {
@@ -54,7 +53,7 @@ impl<'a> TraitsProcessor<'a> {
         };
 
         if let Some(impls) = self.implementations.get(&trait_decl.name) {
-            for impl_ in impls.iter() {
+            for (index, impl_) in impls.iter().enumerate() {
                 let Some(impl_fn) = impl_.function_by_name(&trait_fn.name) else {
                     panic!(
                         "Could not find function {} for {}",
@@ -62,7 +61,7 @@ impl<'a> TraitsProcessor<'a> {
                     );
                 };
 
-                resolved_impl_pos.insert(collected_ref.1.clone(), impl_fn.body.clone());
+                resolved_impl_pos.insert(collected_ref.1.clone(), index);
             }
         }
 
@@ -76,7 +75,7 @@ impl<'a> TraitsProcessor<'a> {
 fn update_reference(
     ref_name: &str,
     c: &mut Reference,
-    resolved_impl_pos: &BTreeMap<Vec<Type>, Box<Expression>>,
+    resolved_impl_pos: &BTreeMap<Vec<Type>, usize>,
 ) {
     if let Reference::Poly(PolynomialReference {
         name,