perf: improve workspace performance (#490)

* perf: improve workspace performance

* fix fmt

core/engine/src/policy/queries/dependency.rs

@@ -21,6 +21,8 @@ pub struct ShallowAnalyses {
     pub per_rule: Vec<RuleShallowAnalysis>,
     pub diagnostics: Vec<Diagnostic>,
     by_block: HashMap<BlockRef, usize>,
+    rules_by_path: HashMap<Arc<str>, std::ops::Range<usize>>,
+    diags_by_path: HashMap<Arc<str>, std::ops::Range<usize>>,
 }
 
 impl ShallowAnalyses {
@@ -29,6 +31,20 @@ impl ShallowAnalyses {
             .get(block_ref)
             .and_then(|&i| self.per_rule.get(i))
     }
+
+    pub fn rules_for(&self, path: &Arc<str>) -> &[RuleShallowAnalysis] {
+        self.rules_by_path
+            .get(path)
+            .map(|r| &self.per_rule[r.clone()])
+            .unwrap_or(&[])
+    }
+
+    pub fn diags_for(&self, path: &Arc<str>) -> &[Diagnostic] {
+        self.diags_by_path
+            .get(path)
+            .map(|r| &self.diagnostics[r.clone()])
+            .unwrap_or(&[])
+    }
 }
 
 #[derive(Debug, Clone)]
@@ -322,11 +338,15 @@ impl Snapshot {
     ) -> ShallowAnalyses {
         let mut per_rule: Vec<RuleShallowAnalysis> = Vec::new();
         let mut diagnostics: Vec<Diagnostic> = Vec::new();
+        let mut rules_by_path: HashMap<Arc<str>, std::ops::Range<usize>> = HashMap::new();
+        let mut diags_by_path: HashMap<Arc<str>, std::ops::Range<usize>> = HashMap::new();
 
         let mut sorted_paths: Vec<&Arc<str>> = all_parsed.keys().collect();
         sorted_paths.sort();
         for path in sorted_paths {
             let p = &all_parsed[path];
+            let rules_start = per_rule.len();
+            let diags_start = diagnostics.len();
 
             for rule in p.policy.rules() {
                 rule.check_single_entity_scope(path, classifier, &mut diagnostics);
@@ -355,6 +375,8 @@ impl Snapshot {
                     .collect()
             });
             per_rule.extend(policy_shallow.iter().cloned());
+            rules_by_path.insert(path.clone(), rules_start..per_rule.len());
+            diags_by_path.insert(path.clone(), diags_start..diagnostics.len());
         }
 
         let by_block = per_rule
@@ -375,6 +397,8 @@ impl Snapshot {
             per_rule,
             diagnostics,
             by_block,
+            rules_by_path,
+            diags_by_path,
         }
     }
 
@@ -534,6 +558,7 @@ impl Snapshot {
         graph: &DependencyGraph,
         order: &[PropertyPath],
         rule_by_ref: &HashMap<BlockRef, Arc<Block>>,
+        shallow: &ShallowAnalyses,
         members: &HashSet<Arc<str>>,
         intellisense: &SharedIntelliSense,
         dictionary_types: SharedDictionaryTypes,
@@ -560,16 +585,24 @@ impl Snapshot {
                 }
             }
         }
-        let mut remaining: Vec<&BlockRef> = rule_by_ref
-            .keys()
-            .filter(|key| members.contains(&key.policy_path) && !analyzed.contains(*key))
-            .collect();
+        let mut remaining: Vec<BlockRef> = Vec::new();
+        for member in members {
+            for s in shallow.rules_for(member) {
+                let key = BlockRef {
+                    policy_path: s.policy_path.clone(),
+                    block_id: s.block_id.clone(),
+                };
+                if analyzed.insert(key.clone()) {
+                    remaining.push(key);
+                }
+            }
+        }
         remaining.sort_by(|a, b| {
             a.policy_path
                 .cmp(&b.policy_path)
                 .then_with(|| a.block_id.cmp(&b.block_id))
         });
-        schedule.extend(remaining.into_iter().map(|key| (key.clone(), false)));
+        schedule.extend(remaining.into_iter().map(|key| (key, false)));
 
         for (key, splice) in schedule {
             let Some(rule) = rule_by_ref.get(&key) else {

core/engine/src/policy/queries/diagnostics.rs

@@ -1,7 +1,6 @@
 use std::sync::Arc;
 
 use ahash::{HashMap, HashMapExt, HashSet};
-use petgraph::algo::tarjan_scc;
 
 use crate::policy::ir::PropertyTypeIr;
 use crate::policy::linter::Linter;
@@ -19,13 +18,7 @@ impl Db {
         }
 
         let shallow = self.shallow();
-        out.extend(
-            shallow
-                .diagnostics
-                .iter()
-                .filter(|d| d.is_in(path))
-                .cloned(),
-        );
+        out.extend(shallow.diags_for(path).iter().cloned());
 
         out.extend(self.graph_diagnostics(path));
 
@@ -67,10 +60,7 @@ impl Db {
         let entity_sources = &unit.entity_sources;
         let classifier = &unit.classifier;
 
-        for rule_analysis in &shallow.per_rule {
-            if !rule_analysis.is_in(target) {
-                continue;
-            }
+        for rule_analysis in shallow.rules_for(target) {
             let mut flagged: HashSet<Arc<str>> = HashSet::default();
             for write in &rule_analysis.writes {
                 let PathRoot::Entity { entity, .. } = classifier.classify(&write.path) else {
@@ -110,10 +100,7 @@ impl Db {
         let classifier = &unit.classifier;
         let rule_index = self.rule_by_ref();
 
-        for rule_analysis in &shallow.per_rule {
-            if !rule_analysis.is_in(target) {
-                continue;
-            }
+        for rule_analysis in shallow.rules_for(target) {
             let block_ref = BlockRef {
                 policy_path: rule_analysis.policy_path.clone(),
                 block_id: rule_analysis.block_id.clone(),
@@ -175,10 +162,9 @@ impl Db {
         let mut first_writer: HashMap<Arc<str>, BlockRef> = HashMap::new();
         let mut all_writes: Vec<(BlockRef, bool, Arc<str>)> = Vec::new();
 
-        for rule in &shallow.per_rule {
-            if !visible.contains(&rule.policy_path) {
-                continue;
-            }
+        let mut sorted_members: Vec<&Arc<str>> = visible.iter().collect();
+        sorted_members.sort();
+        for rule in sorted_members.iter().flat_map(|m| shallow.rules_for(m)) {
             let in_target = rule.is_in(target);
             let block_ref = BlockRef {
                 policy_path: rule.policy_path.clone(),
@@ -340,7 +326,7 @@ impl Db {
         let Some(parsed) = self.parsed(target) else {
             return out;
         };
-        let all_paths: HashSet<Arc<str>> = self.document_paths().into_iter().collect();
+        let all_paths = self.path_set();
 
         for imported in parsed.policy.imports() {
             if !all_paths.contains(imported) {
@@ -352,29 +338,8 @@ impl Db {
             }
         }
 
-        let import_graph = self.import_graph();
-        for scc in tarjan_scc(&import_graph.graph) {
-            let is_cycle = scc.len() > 1
-                || scc
-                    .first()
-                    .is_some_and(|&idx| import_graph.graph.contains_edge(idx, idx));
-            if !is_cycle {
-                continue;
-            }
-            let mut members: Vec<Arc<str>> = scc
-                .iter()
-                .map(|&idx| import_graph.graph[idx].clone())
-                .collect();
-            if !members.iter().any(|p| p == target) {
-                continue;
-            }
-            members.sort();
-            let rendered: Vec<String> = members.iter().map(|p| p.to_string()).collect();
-            out.push(Diagnostic::error(
-                DiagnosticCode::CircularImport,
-                DiagnosticLocation::policy(target.clone()),
-                format!("circular import among: {}", rendered.join(", ")),
-            ));
+        if let Some(cycles) = self.import_cycles().get(target) {
+            out.extend(cycles.iter().cloned());
         }
         out
     }

core/engine/src/workspace/db.rs

@@ -142,6 +142,8 @@ pub struct Snapshot {
     pub(crate) units: RefCell<HashMap<usize, Arc<Unit>>>,
     pub(crate) policy_diagnostics: RefCell<HashMap<Arc<str>, Arc<Vec<Diagnostic>>>>,
     pub(crate) eval_artifacts: RefCell<HashMap<Arc<str>, Arc<EvalArtifact>>>,
+    pub(crate) path_set: OnceCell<Arc<HashSet<Arc<str>>>>,
+    pub(crate) import_cycles: OnceCell<Arc<HashMap<Arc<str>, Vec<Diagnostic>>>>,
 }
 
 pub struct Unit {
@@ -243,6 +245,52 @@ impl Db {
         self.inputs.borrow().documents.keys().cloned().collect()
     }
 
+    pub(crate) fn path_set(&self) -> Arc<HashSet<Arc<str>>> {
+        let snap = self.snapshot();
+        snap.path_set
+            .get_or_init(|| Arc::new(self.document_paths().into_iter().collect()))
+            .clone()
+    }
+
+    pub(crate) fn import_cycles(&self) -> Arc<HashMap<Arc<str>, Vec<Diagnostic>>> {
+        use crate::workspace::types::{DiagnosticCode, DiagnosticLocation};
+        use petgraph::algo::tarjan_scc;
+
+        let snap = self.snapshot();
+        snap.import_cycles
+            .get_or_init(|| {
+                let import_graph = &snap.import_graph;
+                let mut out: HashMap<Arc<str>, Vec<Diagnostic>> = HashMap::new();
+                for scc in tarjan_scc(&import_graph.graph) {
+                    let is_cycle = scc.len() > 1
+                        || scc
+                            .first()
+                            .is_some_and(|&idx| import_graph.graph.contains_edge(idx, idx));
+                    if !is_cycle {
+                        continue;
+                    }
+                    let mut members: Vec<Arc<str>> = scc
+                        .iter()
+                        .map(|&idx| import_graph.graph[idx].clone())
+                        .collect();
+                    members.sort();
+                    let rendered: Vec<String> = members.iter().map(|p| p.to_string()).collect();
+                    let message = format!("circular import among: {}", rendered.join(", "));
+                    for member in &members {
+                        out.entry(member.clone())
+                            .or_default()
+                            .push(Diagnostic::error(
+                                DiagnosticCode::CircularImport,
+                                DiagnosticLocation::policy(member.clone()),
+                                message.clone(),
+                            ));
+                    }
+                }
+                Arc::new(out)
+            })
+            .clone()
+    }
+
     pub fn raw_document(&self, path: &str) -> Option<Arc<DecisionContent>> {
         self.inputs.borrow().documents.get(path).cloned()
     }
@@ -469,10 +517,10 @@ impl Db {
         unit.enriched_once
             .get_or_init(|| {
                 let snap = self.snapshot();
-                let subset: HashMap<Arc<str>, Arc<ParsedPolicy>> = snap
-                    .all_parsed
+                let subset: HashMap<Arc<str>, Arc<ParsedPolicy>> = unit
+                    .members
                     .iter()
-                    .filter(|(p, _)| unit.members.contains(*p))
+                    .filter_map(|m| snap.all_parsed.get_key_value(m))
                     .map(|(p, v)| (p.clone(), v.clone()))
                     .collect();
                 let base_scope = Snapshot::compute_base_scope(&subset, &unit.entity_sources);
@@ -484,6 +532,7 @@ impl Db {
                     &unit.dep_graph,
                     &unit.execution_order,
                     &snap.rule_by_ref,
+                    &snap.shallow,
                     &unit.members,
                     &self.intellisense,
                     Rc::new(unit.dictionary_types()),
@@ -554,23 +603,39 @@ impl Db {
         let opcode_cache = self.opcode_cache_of_unit(&unit);
         let input_schema = self.input_schema(policy);
         let eval_graph = EvalGraph::from_graph(&unit.dep_graph);
-        let reads: HashMap<BlockRef, Arc<[PropertyRead]>> = snap
-            .rule_by_ref
-            .keys()
-            .filter(|r| unit.members.contains(&r.policy_path))
-            .filter_map(|r| {
-                snap.shallow
-                    .for_block(r)
-                    .map(|s| (r.clone(), Arc::from(s.reads.clone())))
+        let reads: HashMap<BlockRef, Arc<[PropertyRead]>> = unit
+            .members
+            .iter()
+            .flat_map(|m| snap.shallow.rules_for(m))
+            .map(|s| {
+                (
+                    BlockRef {
+                        policy_path: s.policy_path.clone(),
+                        block_id: s.block_id.clone(),
+                    },
+                    Arc::from(s.reads.clone()),
+                )
             })
             .collect();
 
         let intellisense = self.intellisense();
         let entity_form = EntityForm::new(unit.entity_sources.as_ref());
-        let read_plans: HashMap<BlockRef, BlockReadPlan> = snap
-            .rule_by_ref
+        let unit_refs: Vec<(&BlockRef, &Arc<Block>)> = unit
+            .members
             .iter()
-            .filter(|(r, _)| unit.members.contains(&r.policy_path))
+            .flat_map(|m| snap.shallow.rules_for(m))
+            .filter_map(|s| {
+                let block_ref = BlockRef {
+                    policy_path: s.policy_path.clone(),
+                    block_id: s.block_id.clone(),
+                };
+                snap.rule_by_ref
+                    .get_key_value(&block_ref)
+                    .map(|(r, b)| (r, b))
+            })
+            .collect();
+        let read_plans: HashMap<BlockRef, BlockReadPlan> = unit_refs
+            .into_iter()
             .map(|(r, block)| {
                 let mut flatten = |src: &Arc<str>, kind: ExpressionKind| -> Vec<Arc<str>> {
                     if src.is_empty() {
@@ -636,10 +701,6 @@ impl Db {
             .cloned()
     }
 
-    pub fn import_graph(&self) -> Arc<ImportGraph> {
-        self.snapshot().import_graph.clone()
-    }
-
     pub fn shallow(&self) -> Arc<ShallowAnalyses> {
         self.snapshot().shallow.clone()
     }
@@ -730,6 +791,8 @@ impl Snapshot {
             units: RefCell::new(HashMap::new()),
             policy_diagnostics: RefCell::new(HashMap::new()),
             eval_artifacts: RefCell::new(HashMap::new()),
+            path_set: OnceCell::new(),
+            import_cycles: OnceCell::new(),
         }
     }
 
@@ -776,9 +839,9 @@ impl Snapshot {
         shallow: &ShallowAnalyses,
     ) -> Unit {
         let member_set: HashSet<Arc<str>> = members.iter().cloned().collect();
-        let subset: HashMap<Arc<str>, Arc<ParsedPolicy>> = all_parsed
+        let subset: HashMap<Arc<str>, Arc<ParsedPolicy>> = members
             .iter()
-            .filter(|(p, _)| member_set.contains(*p))
+            .filter_map(|m| all_parsed.get_key_value(m))
             .map(|(p, v)| (p.clone(), v.clone()))
             .collect();
 
@@ -788,10 +851,11 @@ impl Snapshot {
         let data_model_paths = Self::compute_data_model_paths(&subset);
         let classifier = Self::compute_path_classifier(&subset);
 
-        let per_rule: Vec<&RuleShallowAnalysis> = shallow
-            .per_rule
+        let mut sorted_members: Vec<&Arc<str>> = members.iter().collect();
+        sorted_members.sort();
+        let per_rule: Vec<&RuleShallowAnalysis> = sorted_members
             .iter()
-            .filter(|r| member_set.contains(&r.policy_path))
+            .flat_map(|m| shallow.rules_for(m))
             .collect();
         let dep_graph = Self::compute_graph(&per_rule, &data_model_paths, &entity_sources);
         let execution_order = Self::compute_execution_order(&dep_graph);