localize store usage further

cozo-core/src/query/eval.rs

@@ -52,13 +52,16 @@ impl<'a> SessionTx<'a> {
     pub(crate) fn stratified_magic_evaluate(
         &self,
         strata: &[CompiledProgram],
-        store_lifetimes: BTreeMap<Symbol, usize>,
+        store_lifetimes: BTreeMap<MagicSymbol, usize>,
         total_num_to_take: Option<usize>,
         num_to_skip: Option<usize>,
         poison: Poison,
     ) -> Result<(InMemRelation, bool)> {
         let mut stores = BTreeMap::new();
         let mut early_return = false;
+        let entry_symbol = MagicSymbol::Muggle {
+            inner: Symbol::new(PROG_ENTRY, SourceSpan(0, 0)),
+        };
         for (stratum, cur_prog) in strata.iter().enumerate() {
             debug!("stratum {}", stratum);
             for (rule_name, rule_set) in cur_prog {
@@ -67,24 +70,33 @@ impl<'a> SessionTx<'a> {
 
             early_return = self.semi_naive_magic_evaluate(
                 cur_prog,
-                &stores,
+                &mut stores,
                 total_num_to_take,
                 num_to_skip,
                 poison.clone(),
             )?;
+            // remove stores that have outlived their usefulness!
+            stores = stores
+                .into_iter()
+                .filter(|(name, _)| {
+                    if *name == entry_symbol {
+                        return true;
+                    }
+                    match store_lifetimes.get(name) {
+                        None => false,
+                        Some(n) => *n > stratum,
+                    }
+                })
+                .collect()
         }
-        let ret_area = stores
-            .remove(&MagicSymbol::Muggle {
-                inner: Symbol::new(PROG_ENTRY, SourceSpan(0, 0)),
-            })
-            .ok_or(NoEntryError)?;
+        let ret_area = stores.remove(&entry_symbol).ok_or(NoEntryError)?;
         Ok((ret_area, early_return))
     }
     /// returns true if early return is activated
     fn semi_naive_magic_evaluate(
         &self,
         prog: &CompiledProgram,
-        stores: &BTreeMap<MagicSymbol, InMemRelation>,
+        stores: &mut BTreeMap<MagicSymbol, InMemRelation>,
         total_num_to_take: Option<usize>,
         num_to_skip: Option<usize>,
         poison: Poison,
@@ -202,7 +214,7 @@ impl<'a> SessionTx<'a> {
         &self,
         rule_symb: &MagicSymbol,
         ruleset: &[CompiledRule],
-        stores: &BTreeMap<MagicSymbol, InMemRelation>,
+        stores: &mut BTreeMap<MagicSymbol, InMemRelation>,
         changed: &mut BTreeMap<&MagicSymbol, bool>,
         limiter: &mut QueryLimiter,
         poison: Poison,
@@ -238,7 +250,7 @@ impl<'a> SessionTx<'a> {
         &self,
         rule_symb: &MagicSymbol,
         ruleset: &[CompiledRule],
-        stores: &BTreeMap<MagicSymbol, InMemRelation>,
+        stores: &mut BTreeMap<MagicSymbol, InMemRelation>,
         changed: &mut BTreeMap<&MagicSymbol, bool>,
         poison: Poison,
     ) -> Result<()> {
@@ -281,7 +293,7 @@ impl<'a> SessionTx<'a> {
         &self,
         rule_symb: &MagicSymbol,
         ruleset: &[CompiledRule],
-        stores: &BTreeMap<MagicSymbol, InMemRelation>,
+        stores: &mut BTreeMap<MagicSymbol, InMemRelation>,
         changed: &mut BTreeMap<&MagicSymbol, bool>,
         limiter: &mut QueryLimiter,
         poison: Poison,
@@ -410,7 +422,7 @@ impl<'a> SessionTx<'a> {
         rule_symb: &MagicSymbol,
         ruleset: &[CompiledRule],
         epoch: u32,
-        stores: &BTreeMap<MagicSymbol, InMemRelation>,
+        stores: &mut BTreeMap<MagicSymbol, InMemRelation>,
         prev_changed: &BTreeMap<&MagicSymbol, bool>,
         changed: &mut BTreeMap<&MagicSymbol, bool>,
         limiter: &mut QueryLimiter,
@@ -484,7 +496,7 @@ impl<'a> SessionTx<'a> {
         rule_symb: &MagicSymbol,
         ruleset: &[CompiledRule],
         epoch: u32,
-        stores: &BTreeMap<MagicSymbol, InMemRelation>,
+        stores: &mut BTreeMap<MagicSymbol, InMemRelation>,
         prev_changed: &BTreeMap<&MagicSymbol, bool>,
         changed: &mut BTreeMap<&MagicSymbol, bool>,
         poison: Poison,

cozo-core/src/query/stratify.rs

@@ -13,10 +13,7 @@ use itertools::Itertools;
 use miette::{ensure, Diagnostic, Result};
 use thiserror::Error;
 
-use crate::data::program::{
-    AlgoRuleArg, NormalFormAlgoOrRules, NormalFormAtom, NormalFormProgram,
-    StratifiedNormalFormProgram,
-};
+use crate::data::program::{AlgoRuleArg, MagicSymbol, NormalFormAlgoOrRules, NormalFormAtom, NormalFormProgram, StratifiedNormalFormProgram};
 use crate::data::symb::{Symbol, PROG_ENTRY};
 use crate::parse::SourceSpan;
 use crate::query::graph::{
@@ -215,7 +212,7 @@ fn make_scc_reduced_graph<'a>(
 
 impl NormalFormProgram {
     /// returns the stratified program and the store lifetimes of the intermediate relations
-    pub(crate) fn stratify(self) -> Result<(StratifiedNormalFormProgram, BTreeMap<Symbol, usize>)> {
+    pub(crate) fn stratify(self) -> Result<(StratifiedNormalFormProgram, BTreeMap<MagicSymbol, usize>)> {
         // prerequisite: the program is already in disjunctive normal form
         // 0. build a graph of the program
         let prog_entry: &Symbol = &Symbol::new(PROG_ENTRY, SourceSpan(0, 0));
@@ -262,7 +259,8 @@ impl NormalFormProgram {
                 if let Some(fr_stratum) = invert_sort_result.get(fr_idx) {
                     for (to, _) in tos {
                         let used_in = n_strata - 1 - *fr_stratum;
-                        match store_lifetimes.entry((*to).clone()) {
+                        let magic_to = MagicSymbol::Muggle { inner: (*to).clone() };
+                        match store_lifetimes.entry(magic_to) {
                             Entry::Vacant(e) => {
                                 e.insert(used_in);
                             }