From f10e06ffa63923d57af123b1372d62665ffb15ba Mon Sep 17 00:00:00 2001
From: Stu Hood <stuhood@gmail.com>
Date: Thu, 27 Oct 2022 22:15:26 -0700
Subject: [PATCH] Expose mutable nodes as `@rule(_mutable=True)`

---
 src/python/pants/engine/internals/native_engine.pyi | 1 +
 src/python/pants/engine/internals/scheduler.py      | 1 +
 src/python/pants/engine/rules.py                    | 9 +++++++++
 src/rust/engine/src/externs/interface.rs            | 2 ++
 src/rust/engine/src/nodes.rs                        | 7 +++++++
 src/rust/engine/src/tasks.rs                        | 3 +++
 6 files changed, 23 insertions(+)

diff --git a/src/python/pants/engine/internals/native_engine.pyi b/src/python/pants/engine/internals/native_engine.pyi
index 0a26ba523ae..fdfb23db510 100644
--- a/src/python/pants/engine/internals/native_engine.pyi
+++ b/src/python/pants/engine/internals/native_engine.pyi
@@ -291,6 +291,7 @@ def tasks_task_begin(
     side_effecting: bool,
     engine_aware_return_type: bool,
     cacheable: bool,
+    mutable: bool,
     name: str,
     desc: str,
     level: int,
diff --git a/src/python/pants/engine/internals/scheduler.py b/src/python/pants/engine/internals/scheduler.py
index 7d721d43ba0..08177c036a8 100644
--- a/src/python/pants/engine/internals/scheduler.py
+++ b/src/python/pants/engine/internals/scheduler.py
@@ -643,6 +643,7 @@ def register_task(rule: TaskRule) -> None:
             side_effecting=any(issubclass(t, SideEffecting) for t in rule.input_selectors),
             engine_aware_return_type=issubclass(rule.output_type, EngineAwareReturnType),
             cacheable=rule.cacheable,
+            mutable=rule.mutable,
             name=rule.canonical_name,
             desc=rule.desc or "",
             level=rule.level.level,
diff --git a/src/python/pants/engine/rules.py b/src/python/pants/engine/rules.py
index ffd70fd1142..e93d242c581 100644
--- a/src/python/pants/engine/rules.py
+++ b/src/python/pants/engine/rules.py
@@ -74,6 +74,7 @@ def _make_rule(
     masked_types: Iterable[Type],
     *,
     cacheable: bool,
+    mutable: bool,
     canonical_name: str,
     desc: Optional[str],
     level: LogLevel,
@@ -117,6 +118,7 @@ def wrapper(func):
             desc=desc,
             level=level,
             cacheable=cacheable,
+            mutable=mutable,
         )
 
         return func
@@ -177,6 +179,10 @@ def _ensure_type_annotation(
     # a @rule. Although the type may be in scope for callers, it will not be consumable in the
     # `@rule` which declares the type masked.
     "_masked_types",
+    # Indicates that a rule produces a mutable output. Rules which produce mutable outputs will be
+    # placed in cycles with their dependees, such that whenever the dependee is invalidated, so is
+    # the mutable dependency. This can be used for mutable memoization.
+    "_mutable",
 }
 # We don't want @rule-writers to use 'rule_type' or 'cacheable' as kwargs directly,
 # but rather set them implicitly based on the rule annotation.
@@ -206,6 +212,7 @@ def rule_decorator(func, **kwargs) -> Callable:
 
     rule_type: RuleType = kwargs["rule_type"]
     cacheable: bool = kwargs["cacheable"]
+    mutable: bool = kwargs.get("_mutable", False)
     masked_types: tuple[type, ...] = tuple(kwargs.get("_masked_types", ()))
     param_type_overrides: dict[str, type] = kwargs.get("_param_type_overrides", {})
 
@@ -281,6 +288,7 @@ def rule_decorator(func, **kwargs) -> Callable:
         parameter_types,
         masked_types,
         cacheable=cacheable,
+        mutable=mutable,
         canonical_name=effective_name,
         desc=effective_desc,
         level=effective_level,
@@ -488,6 +496,7 @@ class TaskRule:
     desc: Optional[str] = None
     level: LogLevel = LogLevel.TRACE
     cacheable: bool = True
+    mutable: bool = False
 
     def __str__(self):
         return "(name={}, {}, {!r}, {}, gets={})".format(
diff --git a/src/rust/engine/src/externs/interface.rs b/src/rust/engine/src/externs/interface.rs
index 60b8f94fcea..a92023b0832 100644
--- a/src/rust/engine/src/externs/interface.rs
+++ b/src/rust/engine/src/externs/interface.rs
@@ -1119,6 +1119,7 @@ fn tasks_task_begin(
   side_effecting: bool,
   engine_aware_return_type: bool,
   cacheable: bool,
+  mutable: bool,
   name: String,
   desc: String,
   level: u64,
@@ -1139,6 +1140,7 @@ fn tasks_task_begin(
     arg_types,
     masked_types,
     cacheable,
+    mutable,
     name,
     if desc.is_empty() { None } else { Some(desc) },
     py_level.into(),
diff --git a/src/rust/engine/src/nodes.rs b/src/rust/engine/src/nodes.rs
index 5477c8c662a..d62eb331846 100644
--- a/src/rust/engine/src/nodes.rs
+++ b/src/rust/engine/src/nodes.rs
@@ -1470,6 +1470,13 @@ impl Node for NodeKey {
     }
   }
 
+  fn mutable(&self) -> bool {
+    match self {
+      &NodeKey::Task(ref s) => s.task.mutable,
+      _ => false,
+    }
+  }
+
   fn cacheable_item(&self, output: &NodeOutput) -> bool {
     match (self, output) {
       (NodeKey::ExecuteProcess(ref ep), NodeOutput::ProcessResult(ref process_result)) => {
diff --git a/src/rust/engine/src/tasks.rs b/src/rust/engine/src/tasks.rs
index d419abdd527..3c6f7bc7338 100644
--- a/src/rust/engine/src/tasks.rs
+++ b/src/rust/engine/src/tasks.rs
@@ -151,6 +151,7 @@ pub struct Task {
   pub masked_types: Vec<TypeId>,
   pub func: Function,
   pub cacheable: bool,
+  pub mutable: bool,
   pub display_info: DisplayInfo,
 }
 
@@ -234,6 +235,7 @@ impl Tasks {
     arg_types: Vec<TypeId>,
     masked_types: Vec<TypeId>,
     cacheable: bool,
+    mutable: bool,
     name: String,
     desc: Option<String>,
     level: Level,
@@ -246,6 +248,7 @@ impl Tasks {
 
     self.preparing = Some(Task {
       cacheable,
+      mutable,
       product: return_type,
       side_effecting,
       engine_aware_return_type,