Improve Scope Area and Borrow Checking in Function Definitions and Calls

src/borrow_checker.rs

@@ -43,6 +43,34 @@ impl<'a> BorrowChecker<'a> {
             }
         }
     }
+    /// This function helps to manage the scope of borrow checking.
+    /// It accepts a closure `action` which is executed within a new scope.
+    /// The function automatically handles the scope entering and exiting
+    /// by pushing a new `HashMap` to `self.borrows` and popping it after `action` execution.
+    ///
+    /// # Arguments
+    ///
+    /// * `action` - A closure that encapsulates the actions to be performed within the new scope.
+    ///
+    /// # Type Parameters
+    ///
+    /// * `F` - The type of the closure.
+    /// * `T` - The output type of the closure.
+    fn allocate_scope<F, T>(&mut self, action: F) -> T
+    where
+        F: FnOnce(&mut Self) -> T,
+    {
+        // enter new scope
+        self.borrows.push(HashMap::new());
+
+        // apply action within the scope
+        let result = action(self);
+
+        // exit scope
+        self.borrows.pop();
+
+        result
+    }
 
     fn check_statement(&mut self, stmt: &'a Statement) -> BorrowResult {
         match stmt {
@@ -52,15 +80,9 @@ impl<'a> BorrowChecker<'a> {
                 is_borrowed,
             } => self.check_variable_decl(name, value, *is_borrowed),
             Statement::FunctionDef { name, args, body } => {
-                self.check_function_def(name, args, body)
+                self.allocate_scope(|s| s.check_function_def(name, args, body))
             }
-            Statement::Scope(stmts) => {
-                self.borrows.push(HashMap::new());      // enter new scope
-                let result = self.check(stmts);         // check statements
-                self.borrows.pop();                     // exit scope
-
-                result
-            },
+            Statement::Scope(stmts) => self.allocate_scope(|s| s.check(stmts)),
             Statement::Return(expr) => self.check_return(expr),
             Statement::Expr(expr) => self.check_expression(expr),
         }
@@ -94,7 +116,7 @@ impl<'a> BorrowChecker<'a> {
                         _ => {}
                     }
                     // [NOTE] 2023-06-15
-                    // `BorrowState::ImmutBorrowed` and `BorrowState::Initialized` into the borrows hashmap 
+                    // `BorrowState::ImmutBorrowed` and `BorrowState::Initialized` into the borrows hashmap
                     // for the same variable name. This could cause potential issue like overwrite the borrow state.
                     self.insert_borrow(name, BorrowState::ImmutBorrowed);
 
@@ -159,11 +181,12 @@ impl<'a> BorrowChecker<'a> {
         // check args if exists
         if let Some(args) = args {
             for (arg, is_borrowed) in args {
+                // Insert each argument into the current scope as an initialized variable
+                self.insert_borrow(arg, BorrowState::Initialized);
+
                 if *is_borrowed {
                     self.borrow_imm(arg)?;
                 }
-
-                self.declare(arg.as_str())?;
             }
         }
 
@@ -183,7 +206,6 @@ impl<'a> BorrowChecker<'a> {
             }
 
             scope.insert(var, BorrowState::Uninitialized);
-
             return Ok(());
         }
 
@@ -442,7 +464,10 @@ mod borrow_tests {
         let result = setup(input);
         let result = checker.check(&result);
 
-        assert_eq!(result, Err(BorrowError::VariableNotDefined("b".to_string())));
+        assert_eq!(
+            result,
+            Err(BorrowError::VariableNotDefined("b".to_string()))
+        );
     }
 
     #[test]
@@ -502,7 +527,10 @@ mod borrow_tests {
         let result = setup(input);
         let result = checker.check(&result);
 
-        assert_eq!(result, Err(BorrowError::VariableNotDefined("z".to_string())));
+        assert_eq!(
+            result,
+            Err(BorrowError::VariableNotDefined("z".to_string()))
+        );
     }
 
     #[test]
@@ -539,7 +567,6 @@ mod borrow_tests {
     }
 
     #[test]
-    #[ignore = "function's parameter is not a variable. So it does not need to initialize."]
     fn check_borrow_in_function_decl() {
         let mut checker = BorrowChecker::new();
 
@@ -558,6 +585,43 @@ mod borrow_tests {
         assert_eq!(result, Ok(()));
     }
 
+    #[test]
+    fn check_borrow_multiple_function_decl() {
+        let mut checker = BorrowChecker::new();
+
+        let input = r#"
+            function foo(a) {
+                let b = &a;
+            }
+
+            function bar(a) {
+                let b = &a;
+                let d = foo(&b);
+
+                return d;
+            }
+
+            function baz(a, b) {
+                let c = foo(&a);
+                let d = bar(&b);
+
+                return c + d;
+            }
+
+            let x = 5;
+            let y = 10;
+
+            foo(&x);
+            bar(&y);
+            baz(foo(&x), bar(&y));
+        "#;
+
+        let result = setup(input);
+        let result = checker.check(&result);
+
+        assert_eq!(result, Ok(()));
+    }
+
     #[test]
     fn test_nested_scope_with_uninitialized_variable() {
         let mut checker = BorrowChecker::new();

src/errors.rs

@@ -157,4 +157,4 @@ impl std::fmt::Display for OwnerGraphError {
     }
 }
 
-impl std::error::Error for OwnerGraphError {}
+impl std::error::Error for OwnerGraphError {}

src/main.rs

@@ -4,8 +4,8 @@ mod errors;
 mod eval;
 mod lexer;
 mod lifetime;
+mod ownership;
 mod parser;
 mod token;
-mod ownership;
 
 fn main() {}

src/ownership.rs

@@ -1,14 +1,17 @@
 use std::collections::BTreeMap;
 
-use crate::{ast::{Statement, Expression}, errors::OwnerGraphError};
+use crate::{
+    ast::{Expression, Statement},
+    errors::OwnerGraphError,
+};
 
 /// A structure representing an ownership graph.
 /// The `collections::BTreeMap` preserves the insertion order of the keys.
 #[derive(Debug, PartialEq, Eq, Clone)]
 struct OwnershipGraph {
     /// The `graph` field represents the relationship
     /// between the owner and the variables that it owns.
-    /// 
+    ///
     /// The `key` is a variable that is borrowing, and the `value`
     /// is a vector of variables that are borrowed by the key.
     graph: BTreeMap<String, Vec<String>>,
@@ -23,18 +26,18 @@ impl OwnershipGraph {
     }
 
     /// `OwnershipGraph::add_owner()` adds a new node(`owner`) to the graph.
-    /// 
+    ///
     /// A node is added when a new variable is declared.
     pub fn add_owner(&mut self, owner: &str, variable: &str) {
         let entry = self.graph.entry(owner.to_string()).or_insert_with(Vec::new);
-        
+
         if !entry.contains(&variable.to_string()) {
             entry.push(variable.to_string());
         }
     }
 
     /// `OwnershipGraph::add_borrower()` adds an edge between nodes in the graph.
-    /// 
+    ///
     /// An edge is added when a variable borrows from another variable.
     pub fn add_borrower(&mut self, owner: &str, borrower: &str) {
         if let Some(borrowers) = self.graph.get_mut(owner) {
@@ -45,7 +48,7 @@ impl OwnershipGraph {
     }
 
     /// `OwnershipGraph::remove_owner()` removes a node(`owner`) from the graph.
-    /// 
+    ///
     /// A node is removed when a variable is out of scope.
     pub fn remove_owner(&mut self, owner: &str, variable: &str) {
         if let Some(owners) = self.graph.get_mut(owner) {
@@ -54,7 +57,7 @@ impl OwnershipGraph {
     }
 
     /// `OwnershipGraph::remove_borrower()` removes an edge between nodes in the graph.
-    /// 
+    ///
     /// An edge is removed when a variable is out of scope.
     pub fn remove_borrower(&mut self, owner: &str, borrower: &str) {
         if let Some(borrowers) = self.graph.get_mut(owner) {
@@ -70,7 +73,7 @@ impl Default for OwnershipGraph {
 }
 
 /// `build_ownership_graph()` builds an ownership graph from a list of statements.
-/// 
+///
 /// It iterates through the list of statements (AST generated from the parser) and processes variable
 /// declarations and borrow expressions.
 fn build_ownership_graph(stmts: &[Statement]) -> Result<OwnershipGraph, OwnerGraphError> {
@@ -80,7 +83,11 @@ fn build_ownership_graph(stmts: &[Statement]) -> Result<OwnershipGraph, OwnerGra
 
     for stmt in stmts {
         match stmt {
-            Statement::VariableDecl { name, is_borrowed, value } => {
+            Statement::VariableDecl {
+                name,
+                is_borrowed,
+                value,
+            } => {
                 if *is_borrowed {
                     if let Some(Expression::Reference(ref_var)) = value {
                         current_owner = ref_var.clone();
@@ -92,12 +99,17 @@ fn build_ownership_graph(stmts: &[Statement]) -> Result<OwnershipGraph, OwnerGra
                 }
                 graph.add_owner(&current_owner, name);
             }
-            Statement::Scope(scope) =>  {
+            Statement::Scope(scope) => {
                 let prev_owner = current_owner.clone();
                 let mut declared_in_scope = vec![];
 
                 for inner_stmt in scope {
-                    if let Statement::VariableDecl { name, value, is_borrowed } = inner_stmt {
+                    if let Statement::VariableDecl {
+                        name,
+                        value,
+                        is_borrowed,
+                    } = inner_stmt
+                    {
                         declared_in_scope.push(name.clone());
 
                         if *is_borrowed {
@@ -178,7 +190,7 @@ mod ownership_graph_tests {
                 value: Some(Expression::Reference("b".into())),
             },
         ];
-        
+
         let graph = build_ownership_graph(&stmts).unwrap();
 
         println!("{:?}", graph);
@@ -232,9 +244,9 @@ mod ownership_graph_tests {
             },
         ];
         let graph = build_ownership_graph(&stmts).unwrap();
-        
+
         println!("{:?}", graph);
-        
+
         assert_eq!(
             graph,
             OwnershipGraph {
@@ -308,7 +320,7 @@ mod ownership_graph_tests {
                 value: Some(Expression::Reference("x".into())),
             }]),
         ];
-        
+
         let graph = build_ownership_graph(&stmts).unwrap();
 
         println!("{:?}", graph);
@@ -351,7 +363,7 @@ mod ownership_graph_tests {
                 value: Some(Expression::Reference("x".into())),
             },
         ];
-        
+
         let graph = build_ownership_graph(&stmts).unwrap();
 
         println!("{:?}", graph);
@@ -426,8 +438,14 @@ mod ownership_graph_tests {
             graph,
             OwnershipGraph {
                 graph: vec![
-                    ("global_var".into(), vec!["x".into(), "y".into(), "z".into()]),
-                    ("x".into(), vec!["a".into(), "d".into(), "g".into(), "h".into()]),
+                    (
+                        "global_var".into(),
+                        vec!["x".into(), "y".into(), "z".into()]
+                    ),
+                    (
+                        "x".into(),
+                        vec!["a".into(), "d".into(), "g".into(), "h".into()]
+                    ),
                     ("y".into(), vec!["b".into(), "e".into()]),
                     ("z".into(), vec!["c".into(), "f".into()]),
                     ("a".into(), vec!["d".into()]),
@@ -442,4 +460,4 @@ mod ownership_graph_tests {
             }
         );
     }
-}
+}