change function transpilation to use scope instead of separate map

src/transpile/error.rs

@@ -1,6 +1,6 @@
 //! Errors that can occur during transpilation.
 
-use std::{collections::BTreeMap, fmt::Display};
+use std::{fmt::Display, sync::Arc};
 
 use getset::Getters;
 use itertools::Itertools;
@@ -13,7 +13,10 @@ use crate::{
     semantic::error::{ConflictingFunctionNames, InvalidFunctionArguments, UnexpectedExpression},
 };
 
-use super::FunctionData;
+use super::{
+    variables::{Scope, VariableType},
+    FunctionData,
+};
 
 /// Errors that can occur during transpilation.
 #[allow(clippy::module_name_repetitions, missing_docs)]
@@ -47,20 +50,21 @@ pub struct MissingFunctionDeclaration {
 
 impl MissingFunctionDeclaration {
     #[cfg_attr(not(feature = "shulkerbox"), expect(unused))]
-    pub(super) fn from_context(
+    pub(super) fn from_scope(identifier_span: Span, scope: &Arc<Scope>) -> Self {
-        identifier_span: Span,
-        functions: &BTreeMap<(String, String), FunctionData>,
-    ) -> Self {
         let own_name = identifier_span.str();
-        let own_program_identifier = identifier_span.source_file().identifier();
+        let alternatives = scope
-        let alternatives = functions
+            .get_variables()
+            .read()
+            .unwrap()
             .iter()
-            .filter_map(|((program_identifier, function_name), data)| {
+            .filter_map(|(name, value)| {
-                let normalized_distance =
+                let data = match value.as_ref() {
-                    strsim::normalized_damerau_levenshtein(own_name, function_name);
+                    VariableType::Function { function_data, .. } => function_data,
-                (program_identifier == own_program_identifier
+                    _ => return None,
-                    && (normalized_distance > 0.8
+                };
-                        || strsim::damerau_levenshtein(own_name, function_name) < 3))
+
+                let normalized_distance = strsim::normalized_damerau_levenshtein(own_name, name);
+                (normalized_distance > 0.8 || strsim::damerau_levenshtein(own_name, name) < 3)
                     .then_some((normalized_distance, data))
             })
             .sorted_by(|a, b| a.0.partial_cmp(&b.0).unwrap_or(std::cmp::Ordering::Equal))

src/transpile/transpiler.rs

@@ -4,6 +4,7 @@ use chksum_md5 as md5;
 use std::{
     collections::{BTreeMap, HashMap},
     ops::Deref,
+    sync::{Arc, OnceLock},
 };
 
 use shulkerbox::datapack::{self, Command, Datapack, Execute};
@@ -37,10 +38,10 @@ use super::{
 #[derive(Debug)]
 pub struct Transpiler {
     datapack: shulkerbox::datapack::Datapack,
+    /// Top-level [`Scope`] for each program identifier
+    scopes: BTreeMap<String, Arc<Scope<'static>>>,
     /// Key: (program identifier, function name)
     functions: BTreeMap<(String, String), FunctionData>,
-    /// Key: (program identifier, function name), Value: (function location, public)
-    function_locations: HashMap<(String, String), (String, bool)>,
     /// Key: alias, Value: target
     aliases: HashMap<(String, String), (String, String)>,
 }
@@ -51,8 +52,8 @@ impl Transpiler {
     pub fn new(pack_format: u8) -> Self {
         Self {
             datapack: shulkerbox::datapack::Datapack::new(pack_format),
+            scopes: BTreeMap::new(),
             functions: BTreeMap::new(),
-            function_locations: HashMap::new(),
             aliases: HashMap::new(),
         }
     }
@@ -74,10 +75,18 @@ impl Transpiler {
         handler: &impl Handler<base::Error>,
     ) -> Result<(), TranspileError> {
         tracing::trace!("Transpiling program declarations");
-
-        let scope = Scope::new();
-
         for program in programs {
+            let program_identifier = program
+                .namespace()
+                .span()
+                .source_file()
+                .identifier()
+                .clone();
+            let scope = self
+                .scopes
+                .entry(program_identifier)
+                .or_insert_with(Scope::new)
+                .to_owned();
             self.transpile_program_declarations(program, &scope, handler);
         }
 
@@ -101,6 +110,11 @@ impl Transpiler {
         );
 
         for identifier_span in always_transpile_functions {
+            let scope = self
+                .scopes
+                .entry(identifier_span.source_file().identifier().to_owned())
+                .or_insert_with(Scope::new)
+                .to_owned();
             self.get_or_transpile_function(&identifier_span, None, &scope, handler)?;
         }
 
@@ -111,7 +125,7 @@ impl Transpiler {
     fn transpile_program_declarations(
         &mut self,
         program: &ProgramFile,
-        scope: &Scope,
+        scope: &Arc<Scope>,
         handler: &impl Handler<base::Error>,
     ) {
         let namespace = program.namespace();
@@ -127,7 +141,7 @@ impl Transpiler {
         &mut self,
         declaration: &Declaration,
         namespace: &Namespace,
-        _scope: &Scope,
+        scope: &Arc<Scope>,
         handler: &impl Handler<base::Error>,
     ) {
         let program_identifier = declaration.span().source_file().identifier().clone();
@@ -148,9 +162,7 @@ impl Transpiler {
                         )
                     })
                     .collect();
-                self.functions.insert(
+                let function_data = FunctionData {
-                    (program_identifier, name),
-                    FunctionData {
                     namespace: namespace.namespace_name().str_content().to_string(),
                     identifier_span: identifier_span.clone(),
                     parameters: function
@@ -165,8 +177,16 @@ impl Transpiler {
                     statements,
                     public: function.is_public(),
                     annotations,
+                };
+                scope.set_variable(
+                    &name,
+                    VariableType::Function {
+                        function_data: function_data.clone(),
+                        path: OnceLock::new(),
                     },
                 );
+                self.functions
+                    .insert((program_identifier, name), function_data);
             }
             Declaration::Import(import) => {
                 let path = import.module().str_content();
@@ -220,7 +240,7 @@ impl Transpiler {
         &mut self,
         identifier_span: &Span,
         arguments: Option<&[&Expression]>,
-        scope: &Scope,
+        scope: &Arc<Scope>,
         handler: &impl Handler<base::Error>,
     ) -> TranspileResult<(String, Option<BTreeMap<String, String>>)> {
         let program_identifier = identifier_span.source_file().identifier();
@@ -229,71 +249,57 @@ impl Transpiler {
             identifier_span.str().to_string(),
         );
         let alias_query = self.aliases.get(&program_query).cloned();
-        let already_transpiled = {
+        let already_transpiled = match scope
-            let locations = &self.function_locations;
+            .get_variable(identifier_span.str())
-            locations
+            .expect("called function should be in scope")
-                .get(&program_query)
+            .as_ref()
+        {
+            VariableType::Function { path, .. } => Some(path.get().is_some()),
+            _ => None,
+        }
+        .expect("called variable should be of type function");
+
+        let function_data = scope
+            .get_variable(identifier_span.str())
             .or_else(|| {
                 alias_query
                     .clone()
-                        .and_then(|q| locations.get(&q).filter(|(_, p)| *p))
+                    .and_then(|(alias_program_identifier, alias_function_name)| {
+                        self.scopes
+                            .get(&alias_program_identifier)
+                            .and_then(|s| s.get_variable(&alias_function_name))
                     })
-                .is_some()
-        };
-        if !already_transpiled {
-            tracing::trace!("Function not transpiled yet, transpiling.");
-
-            let function_scope = scope.new_child();
-
-            let statements = {
-                let functions = &self.functions;
-                let function_data = functions
-                    .get(&program_query)
-                    .or_else(|| {
-                        alias_query
-                            .clone()
-                            .and_then(|q| functions.get(&q).filter(|f| f.public))
             })
             .ok_or_else(|| {
                 let error = TranspileError::MissingFunctionDeclaration(
-                            MissingFunctionDeclaration::from_context(
+                    MissingFunctionDeclaration::from_scope(identifier_span.clone(), scope),
-                                identifier_span.clone(),
-                                functions,
-                            ),
                 );
                 handler.receive(error.clone());
                 error
             })?;
 
+        let (function_data, function_path) = match function_data.as_ref() {
+            VariableType::Function {
+                function_data,
+                path,
+            } => (function_data, path),
+            _ => todo!("correctly throw error on wrong type"),
+        };
+
+        if !already_transpiled {
+            tracing::trace!("Function not transpiled yet, transpiling.");
+
+            let function_scope = Scope::with_parent(scope);
+
             for (i, param) in function_data.parameters.iter().enumerate() {
                 function_scope.set_variable(param, VariableType::FunctionArgument { index: i });
             }
 
-                function_data.statements.clone()
+            let statements = function_data.statements.clone();
-            };
 
             let commands =
                 self.transpile_function(&statements, program_identifier, &function_scope, handler)?;
 
-            let functions = &self.functions;
-            let function_data = functions
-                .get(&program_query)
-                .or_else(|| {
-                    alias_query
-                        .clone()
-                        .and_then(|q| functions.get(&q).filter(|f| f.public))
-                })
-                .ok_or_else(|| {
-                    let error = TranspileError::MissingFunctionDeclaration(
-                        MissingFunctionDeclaration::from_context(
-                            identifier_span.clone(),
-                            functions,
-                        ),
-                    );
-                    handler.receive(error.clone());
-                    error
-                })?;
-
             let modified_name = function_data
                 .annotations
                 .get("deobfuscate")
@@ -333,55 +339,30 @@ impl Transpiler {
                 self.datapack.add_load(&function_location);
             }
 
-            self.function_locations.insert(
+            match scope
-                (
+                .get_variable(identifier_span.str())
-                    program_identifier.to_string(),
+                .expect("variable should be in scope if called")
-                    identifier_span.str().to_string(),
+                .as_ref()
-                ),
+            {
-                (function_location, function_data.public),
+                VariableType::Function { path, .. } => {
-            );
+                    path.set(function_location.clone()).unwrap();
+                }
+                _ => todo!("implement error handling")
+            }
         }
 
-        let parameters = {
+        let parameters = function_data.parameters.clone();
-            let function_data = self
+
-                .functions
+        let function_location = function_path
-                .get(&program_query)
+            .get()
-                .or_else(|| {
-                    alias_query
-                        .clone()
-                        .and_then(|q| self.functions.get(&q).filter(|f| f.public))
-                })
             .ok_or_else(|| {
                 let error = TranspileError::MissingFunctionDeclaration(
-                        MissingFunctionDeclaration::from_context(
+                    MissingFunctionDeclaration::from_scope(identifier_span.clone(), scope),
-                            identifier_span.clone(),
-                            &self.functions,
-                        ),
-                    );
-                    handler.receive(error.clone());
-                    error
-                })?;
-
-            function_data.parameters.clone()
-        };
-
-        let function_location = self
-            .function_locations
-            .get(&program_query)
-            .or_else(|| {
-                alias_query.and_then(|q| self.function_locations.get(&q).filter(|(_, p)| *p))
-            })
-            .ok_or_else(|| {
-                let error = TranspileError::MissingFunctionDeclaration(
-                    MissingFunctionDeclaration::from_context(
-                        identifier_span.clone(),
-                        &self.functions,
-                    ),
                 );
                 handler.receive(error.clone());
                 error
             })
-            .map(|(s, _)| s.to_owned())?;
+            .map(|s| s.to_owned())?;
 
         let arg_count = arguments.map(<[&Expression]>::len);
         if arg_count.is_some_and(|arg_count| arg_count != parameters.len()) {
@@ -440,7 +421,7 @@ impl Transpiler {
         &mut self,
         statements: &[Statement],
         program_identifier: &str,
-        scope: &Scope,
+        scope: &Arc<Scope>,
         handler: &impl Handler<base::Error>,
     ) -> TranspileResult<Vec<Command>> {
         let mut errors = Vec::new();
@@ -466,7 +447,7 @@ impl Transpiler {
         &mut self,
         statement: &Statement,
         program_identifier: &str,
-        scope: &Scope,
+        scope: &Arc<Scope>,
         handler: &impl Handler<base::Error>,
     ) -> TranspileResult<Option<Command>> {
         match statement {
@@ -505,7 +486,7 @@ impl Transpiler {
                 unreachable!("Only literal commands are allowed in functions at this time.")
             }
             Statement::ExecuteBlock(execute) => {
-                let child_scope = scope.new_child();
+                let child_scope = Scope::with_parent(scope);
                 self.transpile_execute_block(execute, program_identifier, &child_scope, handler)
             }
             Statement::DocComment(doccomment) => {
@@ -513,7 +494,7 @@ impl Transpiler {
                 Ok(Some(Command::Comment(content.to_string())))
             }
             Statement::Grouping(group) => {
-                let child_scope = scope.new_child();
+                let child_scope = Scope::with_parent(scope);
                 let statements = group.block().statements();
                 let mut errors = Vec::new();
                 let commands = statements
@@ -580,7 +561,7 @@ impl Transpiler {
     fn transpile_function_call(
         &mut self,
         func: &FunctionCall,
-        scope: &Scope,
+        scope: &Arc<Scope>,
         handler: &impl Handler<base::Error>,
     ) -> TranspileResult<Command> {
         let arguments = func
@@ -616,7 +597,7 @@ impl Transpiler {
         &mut self,
         execute: &ExecuteBlock,
         program_identifier: &str,
-        scope: &Scope,
+        scope: &Arc<Scope>,
         handler: &impl Handler<base::Error>,
     ) -> TranspileResult<Option<Command>> {
         self.transpile_execute_block_internal(execute, program_identifier, scope, handler)
@@ -627,7 +608,7 @@ impl Transpiler {
         &mut self,
         execute: &ExecuteBlock,
         program_identifier: &str,
-        scope: &Scope,
+        scope: &Arc<Scope>,
         handler: &impl Handler<base::Error>,
     ) -> TranspileResult<Option<Execute>> {
         match execute {
@@ -715,7 +696,7 @@ impl Transpiler {
         then: Execute,
         el: Option<&Else>,
         program_identifier: &str,
-        scope: &Scope,
+        scope: &Arc<Scope>,
         handler: &impl Handler<base::Error>,
     ) -> TranspileResult<Option<Execute>> {
         let (_, cond) = cond.clone().dissolve();
@@ -767,7 +748,7 @@ impl Transpiler {
         head: &ExecuteBlockHead,
         tail: Option<Execute>,
         program_identifier: &str,
-        scope: &Scope,
+        scope: &Arc<Scope>,
         handler: &impl Handler<base::Error>,
     ) -> TranspileResult<Option<Execute>> {
         Ok(match head {

src/transpile/variables.rs

@@ -1,11 +1,20 @@
 #![expect(unused)]
 
-use std::{collections::HashMap, sync::RwLock};
+use std::{
+    collections::HashMap,
+    sync::{Arc, OnceLock, RwLock},
+};
 
-use super::Transpiler;
+use strum::EnumIs;
 
-#[derive(Debug, Clone)]
+use super::{FunctionData, Transpiler};
+
+#[derive(Debug, Clone, EnumIs)]
 pub enum VariableType {
+    Function {
+        function_data: FunctionData,
+        path: OnceLock<String>,
+    },
     FunctionArgument {
         index: usize,
     },
@@ -35,28 +44,30 @@ pub enum VariableType {
 
 #[derive(Debug, Default)]
 pub struct Scope<'a> {
-    parent: Option<&'a Scope<'a>>,
+    parent: Option<&'a Arc<Self>>,
-    variables: RwLock<HashMap<String, VariableType>>,
+    variables: RwLock<HashMap<String, Arc<VariableType>>>,
 }
 
 impl<'a> Scope<'a> {
-    pub fn new() -> Self {
+    pub fn new() -> Arc<Self> {
-        Self::default()
+        Arc::new(Self::default())
     }
 
-    pub fn new_child(&'a self) -> Self {
+    pub fn with_parent(parent: &'a Arc<Self>) -> Arc<Self> {
-        Self {
+        Arc::new(Self {
-            parent: Some(self),
+            parent: Some(parent),
-            variables: RwLock::new(HashMap::new()),
+            ..Default::default()
-        }
+        })
     }
 
-    pub fn get_variable(&self, name: &str) -> Option<VariableType> {
+    pub fn get_variable(&self, name: &str) -> Option<Arc<VariableType>> {
         let var = self.variables.read().unwrap().get(name).cloned();
         if var.is_some() {
             var
         } else {
-            self.parent.and_then(|parent| parent.get_variable(name))
+            self.parent
+                .as_ref()
+                .and_then(|parent| parent.get_variable(name))
         }
     }
 
@@ -64,11 +75,15 @@ impl<'a> Scope<'a> {
         self.variables
             .write()
             .unwrap()
-            .insert(name.to_string(), var);
+            .insert(name.to_string(), Arc::new(var));
     }
 
-    pub fn get_parent(&self) -> Option<&'a Self> {
+    pub fn get_variables(&self) -> &RwLock<HashMap<String, Arc<VariableType>>> {
-        self.parent
+        &self.variables
+    }
+
+    pub fn get_parent(&self) -> Option<Arc<Self>> {
+        self.parent.map(|v| v.clone())
     }
 }
 
@@ -81,37 +96,39 @@ mod tests {
     #[test]
     fn test_scope() {
         let scope = Scope::new();
-        {
+        scope.set_variable(
-            let mut variables = scope.variables.write().unwrap();
+            "test",
-            variables.insert(
-                "test".to_string(),
             VariableType::Scoreboard {
                 objective: "test".to_string(),
             },
         );
-        }
+        if let Some(var) = scope.get_variable("test") {
-        match scope.get_variable("test") {
+            match var.as_ref() {
-            Some(VariableType::Scoreboard { objective }) => assert_eq!(objective, "test"),
+                VariableType::Scoreboard { objective } => assert_eq!(objective, "test"),
                 _ => panic!("Incorrect Variable"),
             }
+        } else {
+            panic!("Variable missing")
+        }
     }
 
     #[test]
     fn test_parent() {
         let scope = Scope::new();
-        {
+        scope.set_variable(
-            let mut variables = scope.variables.write().unwrap();
+            "test",
-            variables.insert(
-                "test".to_string(),
             VariableType::Scoreboard {
                 objective: "test".to_string(),
             },
         );
-        }
+        let child = Scope::with_parent(&scope);
-        let child = scope.new_child();
+        if let Some(var) = child.get_variable("test") {
-        match child.get_variable("test") {
+            match var.as_ref() {
-            Some(VariableType::Scoreboard { objective }) => assert_eq!(objective, "test"),
+                VariableType::Scoreboard { objective } => assert_eq!(objective, "test"),
                 _ => panic!("Incorrect Variable"),
             }
+        } else {
+            panic!("Variable missing")
+        }
     }
 }