get list of unknowns

nac3core/src/type_check/signature.rs

@@ -2,7 +2,9 @@
 use super::context::TopLevelContext;
 use super::primitives::*;
 use super::typedef::*;
-use rustpython_parser::ast::{ExpressionType, Statement, StatementType, StringGroup, ComprehensionKind};
+use rustpython_parser::ast::{
+    ComprehensionKind, ExpressionType, Statement, StatementType, StringGroup,
+};
 use std::collections::HashMap;
 
 // TODO: fix condition checking, return error message instead of panic...
@@ -49,10 +51,7 @@ fn name_from_expr<'a>(expr: &'a ExpressionType) -> &'a str {
     }
 }
 
-fn type_from_expr<'a>(
+fn type_from_expr<'a>(ctx: &'a TopLevelContext, expr: &'a ExpressionType) -> Result<Type, String> {
-    ctx: &'a TopLevelContext,
-    expr: &'a ExpressionType,
-) -> Result<Type, String> {
     match expr {
         ExpressionType::Identifier { name } => {
             ctx.get_type(name).ok_or_else(|| "no such type".into())
@@ -161,46 +160,50 @@ fn resolve_function<'a>(
     }
 }
 
-fn get_expression_unknowns<'a>(defined: &mut Vec<&'a str>, unknowns: &mut Vec<&'a str>, expr: &'a ExpressionType) {
+fn get_expr_unknowns<'a>(
+    defined: &mut Vec<&'a str>,
+    unknowns: &mut Vec<&'a str>,
+    expr: &'a ExpressionType,
+) {
     match expr {
         ExpressionType::BoolOp { values, .. } => {
             for v in values.iter() {
-                get_expression_unknowns(defined, unknowns, &v.node)
+                get_expr_unknowns(defined, unknowns, &v.node)
             }
         }
         ExpressionType::Binop { a, b, .. } => {
-            get_expression_unknowns(defined, unknowns, &a.node);
+            get_expr_unknowns(defined, unknowns, &a.node);
-            get_expression_unknowns(defined, unknowns, &b.node);
+            get_expr_unknowns(defined, unknowns, &b.node);
         }
         ExpressionType::Subscript { a, b } => {
-            get_expression_unknowns(defined, unknowns, &a.node);
+            get_expr_unknowns(defined, unknowns, &a.node);
-            get_expression_unknowns(defined, unknowns, &b.node);
+            get_expr_unknowns(defined, unknowns, &b.node);
         }
         ExpressionType::Unop { a, .. } => {
-            get_expression_unknowns(defined, unknowns, &a.node);
+            get_expr_unknowns(defined, unknowns, &a.node);
         }
         ExpressionType::Compare { vals, .. } => {
             for v in vals.iter() {
-                get_expression_unknowns(defined, unknowns, &v.node)
+                get_expr_unknowns(defined, unknowns, &v.node)
             }
         }
         ExpressionType::Attribute { value, .. } => {
-            get_expression_unknowns(defined, unknowns, &value.node);
+            get_expr_unknowns(defined, unknowns, &value.node);
         }
         ExpressionType::Call { function, args, .. } => {
-            get_expression_unknowns(defined, unknowns, &function.node);
+            get_expr_unknowns(defined, unknowns, &function.node);
             for v in args.iter() {
-                get_expression_unknowns(defined, unknowns, &v.node)
+                get_expr_unknowns(defined, unknowns, &v.node)
             }
         }
         ExpressionType::List { elements } => {
             for v in elements.iter() {
-                get_expression_unknowns(defined, unknowns, &v.node)
+                get_expr_unknowns(defined, unknowns, &v.node)
             }
         }
         ExpressionType::Tuple { elements } => {
             for v in elements.iter() {
-                get_expression_unknowns(defined, unknowns, &v.node)
+                get_expr_unknowns(defined, unknowns, &v.node)
             }
         }
         ExpressionType::Comprehension { kind, generators } => {
@@ -208,22 +211,22 @@ fn get_expression_unknowns<'a>(defined: &mut Vec<&'a str>, unknowns: &mut Vec<&'
                 unimplemented!()
             }
             let g = &generators[0];
-            get_expression_unknowns(defined, unknowns, &g.iter.node);
+            get_expr_unknowns(defined, unknowns, &g.iter.node);
             let mut scoped = defined.clone();
-            get_expression_unknowns(defined, &mut scoped, &g.target.node);
+            get_expr_unknowns(defined, &mut scoped, &g.target.node);
             for if_expr in g.ifs.iter() {
-                get_expression_unknowns(&mut scoped, unknowns, &if_expr.node);
+                get_expr_unknowns(&mut scoped, unknowns, &if_expr.node);
             }
             match kind.as_ref() {
                 ComprehensionKind::List { element } => {
-                    get_expression_unknowns(&mut scoped, unknowns, &element.node);
+                    get_expr_unknowns(&mut scoped, unknowns, &element.node);
                 }
-                _ => unimplemented!()
+                _ => unimplemented!(),
             }
         }
         ExpressionType::Slice { elements } => {
             for v in elements.iter() {
-                get_expression_unknowns(defined, unknowns, &v.node);
+                get_expr_unknowns(defined, unknowns, &v.node);
             }
         }
         ExpressionType::Identifier { name } => {
@@ -232,14 +235,164 @@ fn get_expression_unknowns<'a>(defined: &mut Vec<&'a str>, unknowns: &mut Vec<&'
             }
         }
         ExpressionType::IfExpression { test, body, orelse } => {
-            get_expression_unknowns(defined, unknowns, &test.node);
+            get_expr_unknowns(defined, unknowns, &test.node);
-            get_expression_unknowns(defined, unknowns, &body.node);
+            get_expr_unknowns(defined, unknowns, &body.node);
-            get_expression_unknowns(defined, unknowns, &orelse.node);
+            get_expr_unknowns(defined, unknowns, &orelse.node);
         }
-        _ => ()
+        _ => (),
     };
 }
 
+struct ExprPattern<'a>(&'a ExpressionType, Vec<usize>, bool);
+
+impl<'a> ExprPattern<'a> {
+    fn new(expr: &'a ExpressionType) -> ExprPattern {
+        let mut pattern = ExprPattern(expr, Vec::new(), true);
+        pattern.find_leaf();
+        pattern
+    }
+
+    fn pointed(&mut self) -> &'a ExpressionType {
+        let mut current = self.0;
+        for v in self.1.iter() {
+            if let ExpressionType::Tuple { elements } = current {
+                current = &elements[*v].node
+            } else {
+                unreachable!()
+            }
+        }
+        current
+    }
+
+    fn find_leaf(&mut self) {
+        let mut current = self.pointed();
+        while let ExpressionType::Tuple { elements } = current {
+            if elements.is_empty() {
+                break;
+            }
+            current = &elements[0].node;
+            self.1.push(0);
+        }
+    }
+
+    fn inc(&mut self) -> bool {
+        loop {
+            if self.1.is_empty() {
+                return false;
+            }
+            let ind = self.1.pop().unwrap() + 1;
+            let parent = self.pointed();
+            if let ExpressionType::Tuple { elements } = parent {
+                if ind < elements.len() {
+                    self.1.push(ind);
+                    self.find_leaf();
+                    return true;
+                }
+            } else {
+                unreachable!()
+            }
+        }
+    }
+}
+
+impl<'a> Iterator for ExprPattern<'a> {
+    type Item = &'a ExpressionType;
+    fn next(&mut self) -> Option<Self::Item> {
+        if self.2 {
+            self.2 = false;
+            Some(self.pointed())
+        } else if self.inc() {
+            Some(self.pointed())
+        } else {
+            None
+        }
+    }
+}
+
+fn get_stmt_unknowns<'a>(
+    defined: &mut Vec<&'a str>,
+    unknowns: &mut Vec<&'a str>,
+    stmts: &'a [Statement],
+) {
+    for stmt in stmts.iter() {
+        match &stmt.node {
+            StatementType::Return { value } => {
+                if let Some(value) = value {
+                    get_expr_unknowns(defined, unknowns, &value.node);
+                }
+            }
+            StatementType::Assign { targets, value } => {
+                get_expr_unknowns(defined, unknowns, &value.node);
+                for target in targets.iter() {
+                    for node in ExprPattern::new(&target.node).into_iter() {
+                        if let ExpressionType::Identifier { name } = node {
+                            let name = name.as_str();
+                            if !defined.contains(&name) {
+                                defined.push(name);
+                            }
+                        } else {
+                            get_expr_unknowns(defined, unknowns, node);
+                        }
+                    }
+                }
+            }
+            StatementType::AugAssign { target, value, .. } => {
+                get_expr_unknowns(defined, unknowns, &target.node);
+                get_expr_unknowns(defined, unknowns, &value.node);
+            }
+            StatementType::AnnAssign { target, value, .. } => {
+                get_expr_unknowns(defined, unknowns, &target.node);
+                if let Some(value) = value {
+                    get_expr_unknowns(defined, unknowns, &value.node);
+                }
+            }
+            StatementType::Expression { expression } => {
+                get_expr_unknowns(defined, unknowns, &expression.node);
+            }
+            StatementType::Global { names } => {
+                for name in names.iter() {
+                    let name = name.as_str();
+                    if !unknowns.contains(&name) {
+                        unknowns.push(name);
+                    }
+                }
+            }
+            StatementType::If { test, body, orelse }
+            | StatementType::While { test, body, orelse } => {
+                get_expr_unknowns(defined, unknowns, &test.node);
+                // we are not very strict at this point...
+                // some identifiers not treated as unknowns may not be resolved
+                // but should be checked during type inference
+                get_stmt_unknowns(defined, unknowns, body.as_slice());
+                if let Some(orelse) = orelse {
+                    get_stmt_unknowns(defined, unknowns, orelse.as_slice());
+                }
+            }
+            StatementType::For { is_async, target, iter, body, orelse } => {
+                if *is_async {
+                    unimplemented!()
+                }
+                get_expr_unknowns(defined, unknowns, &iter.node);
+                for node in ExprPattern::new(&target.node).into_iter() {
+                    if let ExpressionType::Identifier { name } = node {
+                        let name = name.as_str();
+                        if !defined.contains(&name) {
+                            defined.push(name);
+                        }
+                    } else {
+                        get_expr_unknowns(defined, unknowns, node);
+                    }
+                }
+                get_stmt_unknowns(defined, unknowns, body.as_slice());
+                if let Some(orelse) = orelse {
+                    get_stmt_unknowns(defined, unknowns, orelse.as_slice());
+                }
+            }
+            _ => (),
+        }
+    }
+}
+
 pub fn resolve_signatures<'a>(ctx: &mut TopLevelContext<'a>, stmts: &'a [Statement]) {
     for stmt in stmts.iter() {
         match &stmt.node {
@@ -302,7 +455,7 @@ pub fn resolve_signatures<'a>(ctx: &mut TopLevelContext<'a>, stmts: &'a [Stateme
 mod tests {
     use super::*;
     use indoc::indoc;
-    use rustpython_parser::parser::parse_program;
+    use rustpython_parser::parser::{parse_program, parse_statement};
 
     #[test]
     fn test_get_classes() {
@@ -312,14 +465,14 @@ mod tests {
                 b: Test
 
                 def test(self, a: int32) -> Test2:
-                    return b
+                    return self.b
 
             class Bar(Foo, 'FooBar'):
                 def test2(self, a: list[Foo]) -> Test2:
-                    return b
+                    return self.b
 
                 def test3(self, a: list[FooBar2]) -> Test2:
-                    return b
+                    return self.b
         " })
         .unwrap();
         let (mut classes, mut unknowns) = get_typenames(&ast.statements);
@@ -332,4 +485,19 @@ mod tests {
         assert_eq!(&classes, &["Bar", "Foo"]);
         assert_eq!(&unknowns, &["FooBar", "FooBar2", "Test", "Test2"]);
     }
+
+    #[test]
+    fn test_assignment() {
+        let ast = parse_statement(indoc! {"
+            ((a, b), c[i]) = core.foo(x, get_y())
+        " })
+        .unwrap();
+        let mut defined = Vec::new();
+        let mut unknowns = Vec::new();
+        get_stmt_unknowns(&mut defined, &mut unknowns, ast.as_slice());
+        defined.sort();
+        unknowns.sort();
+        assert_eq!(defined.as_slice(), &["a", "b"]);
+        assert_eq!(unknowns.as_slice(), &["c", "core", "get_y", "i", "x"]);
+    }
 }