DeclaredType in TypeChecker

pil-analyzer/src/pil_analyzer.rs

@@ -14,7 +14,7 @@ use powdr_ast::parsed::asm::{
 use powdr_ast::parsed::types::Type;
 use powdr_ast::parsed::visitor::{AllChildren, Children};
 use powdr_ast::parsed::{
-    self, FunctionKind, LambdaExpression, PILFile, PilStatement, SymbolCategory,
+    self, FunctionKind, LambdaExpression, PILFile, PilStatement, SourceReference, SymbolCategory,
     TraitImplementation, TypedExpression,
 };
 use powdr_number::{FieldElement, GoldilocksField};
@@ -29,7 +29,7 @@ use powdr_parser_util::Error;
 
 use crate::traits_resolver::TraitsResolver;
 use crate::type_builtins::constr_function_statement_type;
-use crate::type_inference::infer_types;
+use crate::type_inference::{infer_types, DeclaredType};
 use crate::{side_effect_checker, AnalysisDriver};
 
 use crate::statement_processor::{Counters, PILItem, StatementProcessor};
@@ -303,39 +303,48 @@ impl PILAnalyzer {
                 )
             })
             .flat_map(|(name, (symbol, value))| {
-                let (type_scheme, expr) = match (symbol.kind, value) {
+                let (declared_type, expr) = match (symbol.kind, value) {
                     (SymbolKind::Poly(PolynomialType::Committed), Some(value)) => {
                         // Witness column, move its value (query function) into the expressions to be checked separately.
-                        let type_scheme = type_from_definition(symbol, &None);
 
                         let FunctionValueDefinition::Expression(TypedExpression { e, .. }) = value
                         else {
                             panic!("Invalid value for query function")
                         };
-
+                        let source = e.source_reference().clone();
                         expressions.push((e, query_type.clone().into()));
 
-                        (type_scheme, None)
+                        let declared_type = type_from_definition(symbol, &None)
+                            .map(|ts| ts.into())
+                            .map(|dec: DeclaredType| dec.with_source(source));
+                        (declared_type, None)
                     }
                     (
                         _,
                         Some(FunctionValueDefinition::Expression(TypedExpression {
                             type_scheme,
                             e,
                         })),
-                    ) => (type_scheme.clone(), Some(e)),
+                    ) => {
+                        let source = e.source_reference();
+                        let declared_type = type_scheme
+                            .clone()
+                            .map(|ts| ts.into())
+                            .map(|dec: DeclaredType| dec.with_source(source.clone()));
+                        (declared_type, Some(e))
+                    }
                     (_, value) => {
-                        let type_scheme = type_from_definition(symbol, value);
+                        let declared_type = type_from_definition(symbol, value).map(|ts| ts.into());
 
                         if let Some(FunctionValueDefinition::Array(items)) = value {
                             // Expect all items in the arrays to be field elements.
                             expressions.extend(items.children_mut().map(|e| (e, Type::Fe.into())));
                         }
 
-                        (type_scheme, None)
+                        (declared_type, None)
                     }
                 };
-                Some((name.clone(), (type_scheme, expr)))
+                Some((name.clone(), (declared_type, expr)))
             })
             .collect();
         for expr in &mut self.proof_items {