Add laziness test suite for type checking dependencies (#3219)

crates/pyrefly_util/src/demand_tree.rs

@@ -0,0 +1,152 @@
+/*
+ * Copyright (c) Meta Platforms, Inc. and affiliates.
+ *
+ * This source code is licensed under the MIT license found in the
+ * LICENSE file in the root directory of this source tree.
+ */
+
+//! Demand tree collection for debugging and testing laziness.
+//!
+//! A [`DemandCollector`] records cross-module demand calls into a tree. A
+//! collector is usually owned by a `Transaction`, scoping collection to a
+//! single check run — so parallel checks don't interfere with one another.
+//! Parent/child nesting is tracked via a per-thread scratch stack.
+//!
+//! The tree is machine-readable via serde. `#[serde(default)]` on optional
+//! fields lets the schema grow without breaking existing consumers.
+
+use std::cell::RefCell;
+use std::fmt;
+use std::sync::Arc;
+use std::sync::Mutex;
+
+use serde::Deserialize;
+use serde::Serialize;
+
+/// What kind of cross-module demand a node represents.
+#[derive(Debug, Clone, Serialize, Deserialize)]
+#[serde(tag = "type", rename_all = "snake_case")]
+pub enum DemandKind {
+    /// A leaf event for an Exports-level demand (e.g. `module_exists`,
+    /// `export_exists`). Carries the reason string identifying which
+    /// `LookupExport` method triggered the demand.
+    Exports { reason: String },
+    /// A cross-module Answer lookup span. May have children if the
+    /// computation recursively demanded data from other modules.
+    /// `key` is the `Debug`-formatted lookup key.
+    Answer { key: String },
+}
+
+/// A node in the demand tree.
+#[derive(Debug, Clone, Serialize, Deserialize)]
+pub struct DemandNode {
+    /// The module that made the demand.
+    pub from: String,
+    /// The module the demand was made against.
+    pub target: String,
+    /// What kind of demand this was.
+    pub kind: DemandKind,
+    /// Nested demands made while computing this one.
+    #[serde(default, skip_serializing_if = "Vec::is_empty")]
+    pub children: Vec<DemandNode>,
+}
+
+thread_local! {
+    /// Per-thread stack of in-flight parent spans, used to nest closed demand
+    /// nodes under their caller. This is inherently per-thread scratch space:
+    /// enter and exit always happen on the same thread, and the stack empties
+    /// itself as long as every enter has a matching exit.
+    static STACK: RefCell<Vec<DemandNode>> = const { RefCell::new(Vec::new()) };
+}
+
+/// A demand-tree collection session. Cloning produces another handle to the
+/// same underlying roots (the collector is reference-counted internally).
+#[derive(Clone, Default)]
+pub struct DemandCollector {
+    roots: Arc<Mutex<Vec<DemandNode>>>,
+}
+
+impl DemandCollector {
+    pub fn new() -> Self {
+        Self::default()
+    }
+
+    /// Open a demand span for a cross-module Answer lookup. Hold the returned
+    /// guard for the duration of the demand — dropping the guard (including
+    /// on unwind) closes the span and attaches it to its parent, keeping
+    /// enter/exit balanced even across panics. `key` is the lookup key,
+    /// formatted via `Debug`.
+    #[inline]
+    pub fn enter(
+        &self,
+        from: impl fmt::Display,
+        target: impl fmt::Display,
+        key: impl fmt::Debug,
+    ) -> DemandSpan<'_> {
+        STACK.with(|stack| {
+            stack.borrow_mut().push(DemandNode {
+                from: from.to_string(),
+                target: target.to_string(),
+                kind: DemandKind::Answer {
+                    key: format!("{key:?}"),
+                },
+                children: Vec::new(),
+            });
+        });
+        DemandSpan { collector: self }
+    }
+
+    /// Record a leaf event for an Exports-level demand. `reason` identifies
+    /// which `LookupExport` method triggered the demand.
+    #[inline]
+    pub fn exports_event(&self, from: impl fmt::Display, target: impl fmt::Display, reason: &str) {
+        self.attach(DemandNode {
+            from: from.to_string(),
+            target: target.to_string(),
+            kind: DemandKind::Exports {
+                reason: reason.to_owned(),
+            },
+            children: Vec::new(),
+        });
+    }
+
+    /// Attach a completed node either to the current parent on the stack or,
+    /// if no parent is in flight, to the collector's shared root list.
+    fn attach(&self, node: DemandNode) {
+        let leftover = STACK.with(|stack| {
+            let mut stack = stack.borrow_mut();
+            match stack.last_mut() {
+                Some(parent) => {
+                    parent.children.push(node);
+                    None
+                }
+                None => Some(node),
+            }
+        });
+        if let Some(node) = leftover {
+            self.roots.lock().unwrap().push(node);
+        }
+    }
+
+    /// Take the collected tree roots, leaving the collector empty.
+    pub fn take_roots(&self) -> Vec<DemandNode> {
+        std::mem::take(&mut *self.roots.lock().unwrap())
+    }
+}
+
+/// RAII guard for an in-flight Answer demand span. Dropping the guard —
+/// whether via normal control flow or unwind — pops the span off the
+/// per-thread stack and attaches it to its parent (or to the collector's
+/// roots if no parent is in flight).
+#[must_use = "demand span is closed when the guard is dropped; hold it for the duration of the demand"]
+pub struct DemandSpan<'a> {
+    collector: &'a DemandCollector,
+}
+
+impl Drop for DemandSpan<'_> {
+    fn drop(&mut self) {
+        if let Some(node) = STACK.with(|stack| stack.borrow_mut().pop()) {
+            self.collector.attach(node);
+        }
+    }
+}

crates/pyrefly_util/src/lib.rs

@@ -31,6 +31,7 @@ pub mod absolutize;
 pub mod arc_id;
 pub mod args;
 pub mod assert_size;
+pub mod demand_tree;
 pub mod display;
 pub mod events;
 pub mod forgetter;

pyrefly/Cargo.toml

@@ -1,4 +1,4 @@
-# @generated by autocargo from //pyrefly/pyrefly:[pyrefly,pyrefly_library,pyrefly_lsp_interaction_tests]
+# @generated by autocargo from //pyrefly/pyrefly:[pyrefly,pyrefly_laziness_tests,pyrefly_library,pyrefly_lsp_interaction_tests]
 
 [package]
 name = "pyrefly"
@@ -7,6 +7,7 @@ authors = ["Meta"]
 edition = "2024"
 repository = "https://github.com/facebook/pyrefly"
 license = "MIT"
+build = "test_laziness/build.rs"
 
 [lib]
 path = "lib/lib.rs"
@@ -15,6 +16,10 @@ path = "lib/lib.rs"
 name = "pyrefly"
 path = "bin/main.rs"
 
+[[test]]
+name = "pyrefly_laziness_tests"
+path = "test_laziness/mod.rs"
+
 [[test]]
 name = "pyrefly_lsp_interaction_tests"
 path = "lib/test/lsp/lsp_interaction/mod.rs"

pyrefly/lib/alt/answers.rs

@@ -37,6 +37,27 @@ use crate::alt::traits::Solve;
 use crate::binding::binding::AnyIdx;
 use crate::binding::binding::Exported;
 use crate::binding::binding::Key;
+use crate::binding::binding::KeyAbstractClassCheck;
+use crate::binding::binding::KeyAnnotation;
+use crate::binding::binding::KeyClass;
+use crate::binding::binding::KeyClassBaseType;
+use crate::binding::binding::KeyClassField;
+use crate::binding::binding::KeyClassMetadata;
+use crate::binding::binding::KeyClassMro;
+use crate::binding::binding::KeyClassSynthesizedFields;
+use crate::binding::binding::KeyConsistentOverrideCheck;
+use crate::binding::binding::KeyDecoratedFunction;
+use crate::binding::binding::KeyDecorator;
+use crate::binding::binding::KeyExpect;
+use crate::binding::binding::KeyExport;
+use crate::binding::binding::KeyLegacyTypeParam;
+use crate::binding::binding::KeyTParams;
+use crate::binding::binding::KeyTypeAlias;
+use crate::binding::binding::KeyUndecoratedFunction;
+use crate::binding::binding::KeyVariance;
+use crate::binding::binding::KeyVarianceCheck;
+use crate::binding::binding::KeyYield;
+use crate::binding::binding::KeyYieldFrom;
 use crate::binding::binding::Keyed;
 use crate::binding::bindings::BindingEntry;
 use crate::binding::bindings::BindingTable;

pyrefly/lib/binding/bindings.rs

@@ -52,6 +52,7 @@ use vec1::Vec1;
 use vec1::vec1;
 
 use crate::binding::binding::AnnotationTarget;
+use crate::binding::binding::AnyIdx;
 use crate::binding::binding::Binding;
 use crate::binding::binding::BindingAnnotation;
 use crate::binding::binding::BindingExport;

pyrefly/lib/commands/check.rs

@@ -76,6 +76,7 @@ use crate::state::require::RequireLevels;
 use crate::state::state::State;
 use crate::state::state::Transaction;
 use crate::state::subscriber::ProgressBarSubscriber;
+use crate::state::subscriber::TestSubscriber;
 
 /// Result data from a non-watch check run, used for telemetry logging.
 pub struct CheckResult {
@@ -262,6 +263,10 @@ struct OutputArgs {
     /// Format for pysa report output (json or capnp)
     #[arg(long, value_enum, default_value_t = report::pysa::PysaFormat::Capnp)]
     report_pysa_format: report::pysa::PysaFormat,
+    /// Report the cross-module demand tree (aggregated summary of LookupAnswer
+    /// and LookupExport calls). Useful for analyzing laziness properties.
+    #[arg(long, value_name = "OUTPUT_FILE")]
+    report_demand_tree: Option<PathBuf>,
     /// Generate a CinderX-format type report (experimental, internal-only).
     #[arg(long, value_name = "OUTPUT_DIR", hide = true)]
     report_cinderx: Option<PathBuf>,
@@ -949,15 +954,22 @@ impl CheckArgs {
         }
 
         let type_check_start = Instant::now();
-        let show_progress_bar =
-            self.output.summary != Summary::None && !self.output.no_progress_bar;
-        if show_progress_bar {
-            transaction.set_subscriber(Some(Box::new(ProgressBarSubscriber::new())));
-        }
+        let demand_tree_subscriber = if self.output.report_demand_tree.is_some() {
+            transaction
+                .set_demand_collector(Some(pyrefly_util::demand_tree::DemandCollector::new()));
+            let sub = TestSubscriber::new();
+            transaction.set_subscriber(Some(Box::new(sub.dupe())));
+            Some(sub)
+        } else {
+            let show_progress_bar =
+                self.output.summary != Summary::None && !self.output.no_progress_bar;
+            if show_progress_bar {
+                transaction.set_subscriber(Some(Box::new(ProgressBarSubscriber::new())));
+            }
+            None
+        };
         transaction.run(handles, require, None);
-        if show_progress_bar {
-            transaction.set_subscriber(None);
-        }
+        transaction.set_subscriber(None);
 
         let loads = if self.behavior.check_all {
             transaction.get_all_errors()
@@ -1201,6 +1213,12 @@ impl CheckArgs {
                 report::dependency_graph::dependency_graph(transaction, handles),
             )?;
         }
+        if let Some(path) = &self.output.report_demand_tree {
+            let roots = transaction.take_demand_roots();
+            let module_steps = demand_tree_subscriber.unwrap().finish_detailed();
+            let output = demand_tree_report_json(&roots, &module_steps);
+            fs_anyhow::write(path, output)?;
+        }
         if self.behavior.expectations {
             loads.check_against_expectations()?;
             Ok((CommandExitStatus::Success, output_errors))
@@ -1212,6 +1230,50 @@ impl CheckArgs {
     }
 }
 
+/// Serialize the demand tree and per-module step info as a pretty-printed
+/// JSON document.
+fn demand_tree_report_json(
+    roots: &[pyrefly_util::demand_tree::DemandNode],
+    module_steps: &SmallMap<Handle, crate::state::subscriber::TestModuleInfo>,
+) -> String {
+    use serde::Serialize;
+
+    #[derive(Serialize)]
+    struct ModuleStep {
+        module: String,
+        last_step: &'static str,
+    }
+
+    #[derive(Serialize)]
+    struct Report<'a> {
+        module_steps: Vec<ModuleStep>,
+        demand_tree: &'a [pyrefly_util::demand_tree::DemandNode],
+    }
+
+    let mut steps: Vec<ModuleStep> = module_steps
+        .iter()
+        .map(|(handle, info)| ModuleStep {
+            module: handle.module().as_str().to_owned(),
+            last_step: match info.last_step {
+                Some(crate::state::steps::Step::Load) => "Load",
+                Some(crate::state::steps::Step::Ast) => "Ast",
+                Some(crate::state::steps::Step::Exports) => "Exports",
+                Some(crate::state::steps::Step::Answers) => "Answers",
+                Some(crate::state::steps::Step::Solutions) => "Solutions",
+                None => "Nothing",
+            },
+        })
+        .collect();
+    // Stable ordering makes diffing reports tractable.
+    steps.sort_by(|a, b| a.module.cmp(&b.module));
+
+    let report = Report {
+        module_steps: steps,
+        demand_tree: roots,
+    };
+    serde_json::to_string_pretty(&report).expect("demand tree report should always serialize")
+}
+
 #[cfg(test)]
 mod tests {
     use std::path::PathBuf;