Complexity reduced

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: HashMap<String, Box<Expression>>,
+    pub resolved_impls: BTreeMap<Vec<Type>, Box<Expression>>,
 }
 
 #[derive(

pil-analyzer/src/evaluator.rs

@@ -828,10 +828,9 @@ impl<'a, 'b, T: FieldElement, S: SymbolLookup<'a, T>> Evaluator<'a, 'b, T, S> {
                         ta
                     });
 
-                    let impl_pos = type_args.as_ref().and_then(|type_args| {
-                        let key = type_args.iter().format(",").to_string();
-                        poly.resolved_impls.get(&key).as_ref().copied()
-                    });
+                    let impl_pos = type_args
+                        .as_ref()
+                        .and_then(|type_args| poly.resolved_impls.get(type_args).as_ref().copied());
 
                     match impl_pos {
                         Some(expr) => {
@@ -854,7 +853,7 @@ impl<'a, 'b, T: FieldElement, S: SymbolLookup<'a, T>> Evaluator<'a, 'b, T, S> {
                                 unreachable!()
                             };
 
-                            Value::TraitFunction(Closure {
+                            Value::Closure(Closure {
                                 lambda: body,
                                 environment: vec![],
                                 type_args: local_type_args.clone(),

pil-analyzer/src/expression_processor.rs

@@ -11,7 +11,7 @@ use powdr_ast::{
 use powdr_number::DegreeType;
 use powdr_parser_util::SourceRef;
 use std::{
-    collections::{BTreeSet, HashMap, HashSet},
+    collections::{BTreeMap, BTreeSet, HashMap, HashSet},
     str::FromStr,
 };
 
@@ -394,7 +394,7 @@ impl<'a, D: AnalysisDriver> ExpressionProcessor<'a, D> {
             name: self.driver.resolve_value_ref(&reference.path),
             poly_id: None,
             type_args,
-            resolved_impls: HashMap::new(),
+            resolved_impls: BTreeMap::new(),
         }
     }
 

pil-analyzer/src/pil_analyzer.rs

@@ -18,7 +18,8 @@ use powdr_number::{DegreeType, FieldElement, GoldilocksField};
 
 use powdr_ast::analyzed::{
     type_from_definition, Analyzed, Expression, FunctionValueDefinition, Identity, IdentityKind,
-    PolynomialType, PublicDeclaration, StatementIdentifier, Symbol, SymbolKind, TypedExpression,
+    PolynomialType, PublicDeclaration, Reference, StatementIdentifier, Symbol, SymbolKind,
+    TypedExpression,
 };
 use powdr_parser::{parse, parse_module, parse_type};
 
@@ -327,14 +328,47 @@ impl PILAnalyzer {
     }
 
     pub fn traits_resolution(&mut self) {
-        // TraitsProcessor::new(&mut self.definitions, &self.implementations)
-        //     .traits_resolution(&mut self.identities)
-        TraitsProcessor::new(
-            &mut self.definitions,
-            &mut self.identities,
-            &self.implementations,
-        )
-        .traits_resolution();
+        fn collect_references(expr: &mut Expression) -> Vec<&mut Reference> {
+            match expr {
+                Expression::Reference(_, ref mut r @ Reference::Poly(_)) => vec![r],
+                _ => expr.children_mut().flat_map(collect_references).collect(),
+            }
+        }
+
+        let traits_functions_defs: HashMap<String, FunctionValueDefinition> = self
+            .definitions
+            .iter()
+            .filter_map(|(name, (_, def))| match def {
+                Some(FunctionValueDefinition::TraitFunction(_, _)) => {
+                    Some((name.clone(), def.clone().unwrap()))
+                }
+                _ => None,
+            })
+            .collect();
+
+        let mut references: Vec<&mut Reference> = self
+            .definitions
+            .iter_mut()
+            .flat_map(|(_, (_, def))| match def {
+                Some(FunctionValueDefinition::Expression(TypedExpression { e: expr, .. })) => {
+                    collect_references(expr)
+                }
+                _ => Vec::new(),
+            })
+            .chain(self.identities.iter_mut().flat_map(|identity| {
+                identity
+                    .left
+                    .selector
+                    .iter_mut()
+                    .chain(once(identity.left.expressions.as_mut()))
+                    .chain(identity.right.selector.iter_mut())
+                    .chain(once(identity.right.expressions.as_mut()))
+                    .flat_map(collect_references)
+            }))
+            .collect();
+
+        TraitsProcessor::new(&self.implementations)
+            .traits_resolution(&mut references, traits_functions_defs);
     }
 
     pub fn condense<T: FieldElement>(self) -> Analyzed<T> {