From 7bf6a230f3c8ceec740d463a03cc67879cde5c4d Mon Sep 17 00:00:00 2001
From: Mark Thom <markjordanthom@gmail.com>
Date: Sun, 24 Nov 2019 12:50:53 -0700
Subject: [PATCH] record module-level term_expansion and goal_expansion as
 inner predicates (#228)

---
 src/prolog/forms.rs           |  2 ++
 src/prolog/machine/compile.rs | 20 ++++++++++++---
 src/prolog/machine/modules.rs | 46 +++++++++++++++++++++++++++++++++--
 3 files changed, 62 insertions(+), 6 deletions(-)

diff --git a/src/prolog/forms.rs b/src/prolog/forms.rs
index ed64a67f..c2e4ac14 100644
--- a/src/prolog/forms.rs
+++ b/src/prolog/forms.rs
@@ -337,6 +337,8 @@ pub struct Module {
     pub goal_expansions: (Predicate, VecDeque<TopLevel>),
     pub user_term_expansions: (Predicate, VecDeque<TopLevel>), // term expansions inherited from the user scope.
     pub user_goal_expansions: (Predicate, VecDeque<TopLevel>), // same for goal_expansions.
+    pub local_term_expansions: (Predicate, VecDeque<TopLevel>), // expansions local to the module.
+    pub local_goal_expansions: (Predicate, VecDeque<TopLevel>),
     pub inserted_expansions: bool, // has the module been successfully inserted into toplevel??
     pub is_impromptu_module: bool,
  }
diff --git a/src/prolog/machine/compile.rs b/src/prolog/machine/compile.rs
index 55a916bf..19a20425 100644
--- a/src/prolog/machine/compile.rs
+++ b/src/prolog/machine/compile.rs
@@ -448,7 +448,8 @@ fn add_toplevel_code(wam: &mut Machine, code: Code, indices: IndexStore) {
 }
 
 #[inline]
-fn add_module_code(wam: &mut Machine, mut module: Module, code: Code, indices: IndexStore) {
+fn add_module_code(wam: &mut Machine, mut module: Module, code: Code, indices: IndexStore)
+{
     module.code_dir.extend(indices.code_dir);
     module.op_dir.extend(indices.op_dir.into_iter());
 
@@ -683,7 +684,8 @@ impl ListingCompiler {
                 .or_insert((Predicate::new(), VecDeque::from(vec![])));
 
             if let Some(ref mut module) = &mut self.module {
-                module.add_module_expansion_record(hook, clause.clone(), queue.clone());
+                module.add_expansion_record(hook, clause.clone(), queue.clone());
+                module.add_local_expansion(hook, clause.clone(), queue.clone());
             }
 
             (module_preds.0).0.push(clause);
@@ -918,7 +920,7 @@ fn compile_work_impl(
     mut indices: IndexStore,
     mut results: GatherResult,
 ) -> Result<(), SessionError> {
-    let module_code = compiler.generate_code(
+    let mut module_code = compiler.generate_code(
         results.worker_results,
         wam,
         &mut indices.code_dir,
@@ -969,7 +971,17 @@ fn compile_work_impl(
 
             wam.indices.use_module(&mut wam.code_repo, wam.machine_st.flags, &module)?;
             wam.indices.insert_module(module);
-        } else {
+        } else {                  
+            // compile the module-level goal and term expansions and store
+            // their locations to the module's code_dir.
+            let offset = module_code.len() + toplvl_code.len();
+            let decls = module.take_local_expansions();
+            
+            if !decls.is_empty() {
+                let code = compiler.generate_code(decls, &wam, &mut indices.code_dir, offset)?;
+                module_code.extend(code.into_iter());
+            }
+
             add_module_code(wam, module, module_code, indices);
         }
 
diff --git a/src/prolog/machine/modules.rs b/src/prolog/machine/modules.rs
index ace04f65..b6c07e8f 100644
--- a/src/prolog/machine/modules.rs
+++ b/src/prolog/machine/modules.rs
@@ -7,6 +7,7 @@ use crate::prolog::machine::machine_errors::*;
 use crate::prolog::machine::machine_indices::*;
 
 use std::collections::VecDeque;
+use std::mem;
 
 // Module's and related types are defined in forms.
 impl Module {
@@ -19,6 +20,8 @@ impl Module {
             user_goal_expansions: (Predicate::new(), VecDeque::from(vec![])),
             term_expansions: (Predicate::new(), VecDeque::from(vec![])),
             goal_expansions: (Predicate::new(), VecDeque::from(vec![])),
+            local_term_expansions: (Predicate::new(), VecDeque::from(vec![])),
+            local_goal_expansions: (Predicate::new(), VecDeque::from(vec![])),
             code_dir: CodeDir::new(),
             op_dir: default_op_dir(),
             inserted_expansions: false,
@@ -61,7 +64,7 @@ impl Module {
         Ok(())
     }
 
-    pub fn add_module_expansion_record(
+    pub fn add_expansion_record(
         &mut self,
         hook: CompileTimeHook,
         clause: PredicateClause,
@@ -78,6 +81,45 @@ impl Module {
             }
         }
     }
+
+    pub fn add_local_expansion(
+        &mut self,
+        hook: CompileTimeHook,
+        clause: PredicateClause,
+        queue: VecDeque<TopLevel>,
+    ) {
+        match hook {
+            CompileTimeHook::TermExpansion => {
+                (self.local_term_expansions.0).0.push(clause);
+                self.local_term_expansions.1.extend(queue.into_iter());
+            }
+            CompileTimeHook::GoalExpansion => {
+                (self.local_goal_expansions.0).0.push(clause);
+                self.local_goal_expansions.1.extend(queue.into_iter());
+            }
+            _ => {}
+        }
+    }
+
+    pub fn take_local_expansions(&mut self) -> Vec<(Predicate, VecDeque<TopLevel>)>
+    {
+        let term_expansions =
+            mem::replace(&mut self.local_term_expansions, (Predicate::new(), VecDeque::new()));
+        let goal_expansions =
+            mem::replace(&mut self.local_goal_expansions, (Predicate::new(), VecDeque::new()));
+
+        let mut result = vec![];
+
+        if !(term_expansions.0).0.is_empty() {
+            result.push(term_expansions);
+        }
+
+        if !(goal_expansions.0).0.is_empty() {
+            result.push(goal_expansions);
+        }
+
+        result
+    }
 }
 
 pub trait SubModuleUser {
@@ -283,7 +325,7 @@ impl SubModuleUser for Module {
         self.user_goal_expansions
             .1
             .extend(submodule.goal_expansions.1.iter().cloned());
-	
+
         Ok(())
     }
 }
-- 
2.54.0