From 0f4c7e0c0d685c182f20ee182c9b02f9fac5b16b Mon Sep 17 00:00:00 2001
From: Faith Ekstrand <faith.ekstrand@collabora.com>
Date: Thu, 27 Feb 2025 13:26:05 -0600
Subject: [PATCH] nak: Add a new ConstTracker struct

Reviewed-by: Mel Henning <mhenning@darkrefraction.com>
Part-of: <https://gitlab.freedesktop.org/mesa/mesa/-/merge_requests/33785>
---
 src/nouveau/compiler/nak/const_tracker.rs | 59 +++++++++++++++++++++++
 src/nouveau/compiler/nak/lib.rs           |  1 +
 2 files changed, 60 insertions(+)
 create mode 100644 src/nouveau/compiler/nak/const_tracker.rs

diff --git a/src/nouveau/compiler/nak/const_tracker.rs b/src/nouveau/compiler/nak/const_tracker.rs
new file mode 100644
index 00000000000..629a89b0ee7
--- /dev/null
+++ b/src/nouveau/compiler/nak/const_tracker.rs
@@ -0,0 +1,59 @@
+// Copyright © 2025 Collabora, Ltd.
+// SPDX-License-Identifier: MIT
+
+use crate::ir::*;
+
+use std::collections::HashMap;
+
+pub struct ConstTracker {
+    map: HashMap<SSAValue, SrcRef>,
+}
+
+/// A tracker struct for finding re-materializable constants
+///
+/// Anything which is an immediate, Zero, or a bound cbuf can trivially be
+/// re-materialized anywhere in the shader and it's probably cheaper to do so
+/// than to try and keep them around in GPRs forever.  This is just a helper
+/// struct for implementing this logic in compiler passes.
+impl ConstTracker {
+    pub fn new() -> Self {
+        ConstTracker {
+            map: HashMap::new(),
+        }
+    }
+
+    /// Registers a copy instruction
+    ///
+    /// If the source of the copy is a constant, the destination SSA value and
+    /// the constant value get stored as a key/value pair.
+    pub fn add_copy(&mut self, op: &OpCopy) {
+        let Some(dst) = op.dst.as_ssa() else {
+            return;
+        };
+        debug_assert!(dst.comps() == 1);
+        let dst = dst[0];
+
+        debug_assert!(op.src.src_mod.is_none());
+        let is_const = match &op.src.src_ref {
+            SrcRef::Zero | SrcRef::True | SrcRef::False | SrcRef::Imm32(_) => {
+                true
+            }
+            SrcRef::CBuf(cb) => matches!(cb.buf, CBuf::Binding(_)),
+            _ => false,
+        };
+
+        if is_const {
+            self.map.insert(dst, op.src.src_ref);
+        }
+    }
+
+    /// Tests if the ConstTracker contains the given SSAValue
+    pub fn contains(&self, ssa: &SSAValue) -> bool {
+        self.map.contains_key(ssa)
+    }
+
+    /// Returns the SrcRef associated with this SSAValue, if any
+    pub fn get(&self, ssa: &SSAValue) -> Option<&SrcRef> {
+        self.map.get(ssa)
+    }
+}
diff --git a/src/nouveau/compiler/nak/lib.rs b/src/nouveau/compiler/nak/lib.rs
index 178067d56ba..97b546349cd 100644
--- a/src/nouveau/compiler/nak/lib.rs
+++ b/src/nouveau/compiler/nak/lib.rs
@@ -5,6 +5,7 @@ mod api;
 mod assign_regs;
 mod builder;
 mod calc_instr_deps;
+mod const_tracker;
 mod from_nir;
 mod ir;
 mod legalize;