perf: reuse semantic state for ToolService queries

crates/wright-analyzer/src/canonical/service.rs

@@ -1,4 +1,4 @@
-use std::sync::Arc;
+use std::{ops::Deref, sync::Arc};
 
 use serde_json::{Value as JsonValue, json};
 use workshop_rs::source::{FileId, Span};
@@ -7,14 +7,37 @@ use workshop_rs::{Event, Program};
 use super::analysis::Finding;
 use super::cfg::cfg_response;
 use super::facts::persistent_objects;
-use super::symbols::{ReferenceKind, RuleId, SemanticIndex, Symbol, SymbolId};
+use super::symbols::{Reference, ReferenceKind, RuleId, SemanticIndex, Symbol, SymbolId};
 use crate::analysis::Boundedness;
 use crate::registry::{LintConfig, SkippedRule};
 use crate::service::{ErrorInfo, Origin, Request, Response};
 
+#[derive(Clone)]
+enum ProgramSource<'a> {
+    Borrowed(&'a Program),
+    Shared(Arc<Program>),
+}
+
+impl ProgramSource<'_> {
+    fn as_ref(&self) -> &Program {
+        match self {
+            Self::Borrowed(program) => program,
+            Self::Shared(program) => program,
+        }
+    }
+}
+
+impl Deref for ProgramSource<'_> {
+    type Target = Program;
+
+    fn deref(&self) -> &Self::Target {
+        self.as_ref()
+    }
+}
+
 pub struct SemanticService<'a> {
-    program: &'a Program,
-    index: SemanticIndex,
+    program: ProgramSource<'a>,
+    index: Arc<SemanticIndex>,
     findings: Vec<Finding>,
     skipped: Vec<SkippedRule>,
     origin: Origin,
@@ -62,8 +85,36 @@ impl<'a> SemanticService<'a> {
         config: LintConfig,
         registry: Arc<crate::registry::LintRegistry>,
     ) -> Self {
-        let index = SemanticIndex::build(program);
-        let report = registry.run_report(program, &config);
+        Self::with_program_and_origin_and_config_and_registry(
+            ProgramSource::Borrowed(program),
+            origin,
+            config,
+            registry,
+        )
+    }
+
+    pub fn with_shared_program(
+        program: Arc<Program>,
+        origin: Origin,
+        config: LintConfig,
+        registry: Arc<crate::registry::LintRegistry>,
+    ) -> SemanticService<'static> {
+        SemanticService::<'static>::with_program_and_origin_and_config_and_registry(
+            ProgramSource::Shared(program),
+            origin,
+            config,
+            registry,
+        )
+    }
+
+    fn with_program_and_origin_and_config_and_registry(
+        program: ProgramSource<'a>,
+        origin: Origin,
+        config: LintConfig,
+        registry: Arc<crate::registry::LintRegistry>,
+    ) -> Self {
+        let index = Arc::new(SemanticIndex::build(program.as_ref()));
+        let report = registry.run_report(program.as_ref(), &config);
         Self {
             program,
             index,
@@ -74,6 +125,36 @@ impl<'a> SemanticService<'a> {
             registry,
         }
     }
+
+    pub fn with_lint_config(&self, config: LintConfig) -> Self {
+        let report = self.registry.run_report(self.program.as_ref(), &config);
+        Self {
+            program: self.program.clone(),
+            index: Arc::clone(&self.index),
+            findings: report.findings,
+            skipped: report.skipped,
+            origin: self.origin.clone(),
+            config,
+            registry: Arc::clone(&self.registry),
+        }
+    }
+
+    pub fn references_for_all_symbols(&self) -> JsonValue {
+        JsonValue::Array(
+            self.index
+                .references_for_all_symbols()
+                .into_iter()
+                .map(|references| {
+                    json!(
+                        references
+                            .into_iter()
+                            .map(reference_json)
+                            .collect::<Vec<_>>()
+                    )
+                })
+                .collect(),
+        )
+    }
     pub fn handle_json(&self, request_json: &str) -> String {
         let request: Request = match serde_json::from_str(request_json) {
             Ok(req) => req,
@@ -93,16 +174,16 @@ impl<'a> SemanticService<'a> {
     pub fn handle(&self, request: &Request) -> Response {
         match request {
             Request::Version => Response::Ok { result: json!({"name": "wright-tool", "version": env!("CARGO_PKG_VERSION"), "capabilities": ["program", "rules", "symbols", "references", "usage", "cfg", "findings", "persistentObjects", "lintRules"]}) },
-            Request::Program => Response::Ok { result: json!({"origin": self.origin, "files": file_count(self.program), "globalVariables": self.program.global_variables.len(), "playerVariables": self.program.player_variables.len(), "subroutines": self.program.subroutines.len(), "rules": self.program.rules.len(), "findings": self.findings.len()}) },
+            Request::Program => Response::Ok { result: json!({"origin": self.origin, "files": file_count(self.program.as_ref()), "globalVariables": self.program.global_variables.len(), "playerVariables": self.program.player_variables.len(), "subroutines": self.program.subroutines.len(), "rules": self.program.rules.len(), "findings": self.findings.len()}) },
             Request::ListRules => Response::Ok { result: json!(self.program.rules.iter().enumerate().map(|(id, rule)| json!({"id": id, "name": rule.name, "span": span_json(self.program.rule_span(id))})).collect::<Vec<_>>()) },
             Request::GetRule { rule } => self.rule(*rule as usize),
             Request::ListSymbols { kind } => Response::Ok { result: json!(self.index.symbols().filter(|symbol| kind.as_deref().is_none_or(|kind| symbol.kind.as_str() == kind)).map(symbol_json).collect::<Vec<_>>()) },
             Request::GetSymbol { symbol } => self.index.symbol(SymbolId::from_index(*symbol as usize)).map_or_else(|| self.error("invalid-id", format!("unknown symbol {symbol}")), |symbol| Response::Ok { result: symbol_json(symbol) }),
-            Request::FindReferences { symbol } => { let id = SymbolId::from_index(*symbol as usize); if self.index.symbol(id).is_none() { self.error("invalid-id", format!("unknown symbol {symbol}")) } else { Response::Ok { result: json!(self.index.references(id).into_iter().map(|reference| json!({"kind": reference_kind_name(reference.kind), "span": span_json(reference.span), "rule": reference.rule, "action": reference.action, "value": reference.value})).collect::<Vec<_>>()) } } }
+            Request::FindReferences { symbol } => { let id = SymbolId::from_index(*symbol as usize); if self.index.symbol(id).is_none() { self.error("invalid-id", format!("unknown symbol {symbol}")) } else { Response::Ok { result: json!(self.index.references(id).into_iter().map(reference_json).collect::<Vec<_>>()) } } }
             Request::GetUsage { symbol } => { let id = SymbolId::from_index(*symbol as usize); self.index.symbol(id).map_or_else(|| self.error("invalid-id", format!("unknown symbol {symbol}")), |data| { let usage = self.index.usage(id); Response::Ok { result: json!({"symbol": data.name, "reads": usage.reads, "writes": usage.writes, "calls": usage.calls, "rules": usage.rules}) } }) }
-            Request::GetCfg { rule } => cfg_response(self.program, *rule as usize),
+            Request::GetCfg { rule } => cfg_response(self.program.as_ref(), *rule as usize),
             Request::GetFindings => Response::Ok { result: json!(self.findings.iter().map(finding_json).collect::<Vec<_>>()) },
-            Request::GetPersistentObjects => Response::Ok { result: json!(persistent_objects(self.program)) },
+            Request::GetPersistentObjects => Response::Ok { result: json!(persistent_objects(self.program.as_ref())) },
             Request::LintRules => Response::Ok {
                 result: lint_rules(&self.registry, &self.config, &self.skipped),
             }
@@ -194,6 +275,16 @@ fn reference_kind_name(kind: ReferenceKind) -> &'static str {
         ReferenceKind::Call => "call",
     }
 }
+
+fn reference_json(reference: &Reference) -> JsonValue {
+    json!({
+        "kind": reference_kind_name(reference.kind),
+        "span": span_json(reference.span),
+        "rule": reference.rule,
+        "action": reference.action,
+        "value": reference.value,
+    })
+}
 fn event_name(event: &Event) -> String {
     match event {
         Event::Subroutine(name) => format!("subroutine:{name}"),

crates/wright-analyzer/src/canonical/symbols.rs

@@ -293,6 +293,17 @@ impl SemanticIndex {
             .filter(|reference| reference.symbol == symbol)
             .collect()
     }
+    pub(super) fn references_for_all_symbols(&self) -> Vec<Vec<&Reference>> {
+        let mut grouped = (0..self.symbols.len())
+            .map(|_| Vec::new())
+            .collect::<Vec<_>>();
+        for reference in &self.references {
+            if let Some(references) = grouped.get_mut(reference.symbol.index()) {
+                references.push(reference);
+            }
+        }
+        grouped
+    }
     pub fn usage(&self, symbol: SymbolId) -> UsageSummary {
         let mut usage = UsageSummary::default();
         let mut rules = HashSet::new();

crates/wright-bench/src/main.rs

@@ -6,12 +6,14 @@ use std::time::{Duration, Instant};
 
 use wright_driver::CompilerSession;
 use wright_driver::Profile;
-use wright_driver::config::{SessionConfig, SourceKind};
+use wright_driver::config::{InputSpec, SessionConfig, SourceKind};
+use wright_driver::service::{ToolRequest, ToolResponse, ToolService};
 
 const BENCH_CONTRACT: &str = "wright-bench/v1";
 const BENCH_CONFIG: &str = r#"{
   "iterations": 5,
   "warmup": 1,
+  "semanticQueryRepeats": 5,
   "thresholds": {
     "maxMeanLatencyMs": 500.0,
     "maxEmittedBytes": 200000,
@@ -76,9 +78,24 @@ struct Report {
     iterations: usize,
     thresholds: serde_json::Value,
     fixtures: Vec<FixtureReport>,
+    semantic_queries: SemanticQueryReport,
     summary: serde_json::Value,
 }
 
+#[derive(serde::Serialize)]
+struct SemanticQueryFixtureReport {
+    fixture: &'static str,
+    service_initialization_mean_ms: f64,
+    repeated_queries_mean_ms: f64,
+    loaded_service_mean_ms: f64,
+}
+
+#[derive(serde::Serialize)]
+struct SemanticQueryReport {
+    workload: serde_json::Value,
+    fixtures: Vec<SemanticQueryFixtureReport>,
+}
+
 fn run() -> Result<bool, String> {
     let config: serde_json::Value =
         serde_json::from_str(BENCH_CONFIG).map_err(|e| format!("invalid bench config: {e}"))?;
@@ -140,6 +157,12 @@ fn run() -> Result<bool, String> {
         reports.push(report);
     }
 
+    let semantic_queries = benchmark_semantic_queries(
+        iterations,
+        warmup,
+        config["semanticQueryRepeats"].as_u64().unwrap_or(5) as usize,
+    )?;
+
     let rss_mb = peak_rss_mb();
     if rss_mb
         > config["thresholds"]["maxRssMb"]
@@ -156,6 +179,7 @@ fn run() -> Result<bool, String> {
         iterations,
         thresholds: config["thresholds"].clone(),
         fixtures: reports,
+        semantic_queries,
         summary: serde_json::json!({ "peakRssMb": rss_mb, "regressions": regressions }),
     };
     let out = workspace_root().join("target");
@@ -169,6 +193,122 @@ fn run() -> Result<bool, String> {
     Ok(regressions.is_empty())
 }
 
+fn benchmark_semantic_queries(
+    iterations: usize,
+    warmup: usize,
+    repeats: usize,
+) -> Result<SemanticQueryReport, String> {
+    let root = workspace_root().join("tests/fixtures/workshop");
+    let cases = benchmark_cases()
+        .into_iter()
+        .filter(|(id, _)| {
+            matches!(
+                *id,
+                "synthetic/declarations-rules"
+                    | "synthetic/control-flow"
+                    | "real-world/overpy-cake"
+            )
+        })
+        .collect::<Vec<_>>();
+    let mut fixtures = Vec::new();
+
+    for (id, path) in cases {
+        let source = std::fs::read_to_string(&path)
+            .map_err(|e| format!("cannot read semantic benchmark fixture '{id}': {e}"))?;
+        for _ in 0..warmup {
+            semantic_query_trial(&source, id, &root, repeats)?;
+        }
+
+        let mut initialization = Vec::with_capacity(iterations);
+        let mut queries = Vec::with_capacity(iterations);
+        for _ in 0..iterations {
+            let (service_initialization, repeated_queries) =
+                semantic_query_trial(&source, id, &root, repeats)?;
+            initialization.push(service_initialization);
+            queries.push(repeated_queries);
+        }
+        let initialization_ms = mean(&initialization) * 1000.0;
+        let queries_ms = mean(&queries) * 1000.0;
+        fixtures.push(SemanticQueryFixtureReport {
+            fixture: id,
+            service_initialization_mean_ms: initialization_ms,
+            repeated_queries_mean_ms: queries_ms,
+            loaded_service_mean_ms: initialization_ms + queries_ms,
+        });
+    }
+
+    Ok(SemanticQueryReport {
+        workload: serde_json::json!({
+            "fixtures": fixtures.iter().map(|fixture| fixture.fixture).collect::<Vec<_>>(),
+            "requests": ["rules", "symbols", "references(symbol:0)", "usage(symbol:0)", "cfg(rule:0)", "findings", "persistentObjects", "lint", "lintRules", "inspect"],
+            "repeatsPerFixturePerIteration": repeats,
+            "timing": "CompilerSession::load is completed before timing; initialization measures ToolService::new, repeatedQueries measures the listed request sequence, and loadedService is their sum.",
+        }),
+        fixtures,
+    })
+}
+
+fn semantic_query_trial(
+    source: &str,
+    fixture: &str,
+    root: &Path,
+    repeats: usize,
+) -> Result<(Duration, Duration), String> {
+    let safe_name = fixture.replace('/', "-");
+    let input_dir = workspace_root().join("target/wright-bench-inputs");
+    std::fs::create_dir_all(&input_dir).map_err(|e| e.to_string())?;
+    let path = input_dir.join(format!(
+        "wright-bench-{}-semantic-{safe_name}.ws",
+        std::process::id()
+    ));
+    std::fs::write(&path, source).map_err(|e| e.to_string())?;
+    let result = (|| {
+        let mut session = CompilerSession::new(SessionConfig {
+            input: InputSpec::Path(path.clone()),
+            kind: SourceKind::Workshop,
+            root: Some(root.to_path_buf()),
+            profile: Profile::Compat,
+            ..SessionConfig::default()
+        })
+        .map_err(|e| e.message)?;
+        session.load().map_err(|e| e.message)?;
+
+        let start = Instant::now();
+        let mut service = ToolService::new(&mut session).map_err(|e| e.message)?;
+        let initialization = start.elapsed();
+
+        let requests = [
+            ToolRequest::Rules,
+            ToolRequest::Symbols { kind: None },
+            ToolRequest::References { symbol: 0 },
+            ToolRequest::Usage { symbol: 0 },
+            ToolRequest::Cfg { rule: 0 },
+            ToolRequest::Findings,
+            ToolRequest::PersistentObjects,
+            ToolRequest::Lint,
+            ToolRequest::LintRules,
+        ];
+        let start = Instant::now();
+        for _ in 0..repeats {
+            for request in &requests {
+                if let ToolResponse::Error { error } = service.handle(request) {
+                    return Err(format!("{fixture}: {error:?}"));
+                }
+            }
+            let inspection = service.inspect();
+            if !inspection.ok {
+                return Err(format!(
+                    "{fixture}: inspect failed: {:?}",
+                    inspection.diagnostics
+                ));
+            }
+        }
+        Ok((initialization, start.elapsed()))
+    })();
+    let _ = std::fs::remove_file(&path);
+    result
+}
+
 fn compile(
     source: &str,
     fixture: &str,