From 6293f7f49e5137d0c782ef757d97f9228cd7ea74 Mon Sep 17 00:00:00 2001
From: Alex Koshelev <koshelev@meta.com>
Date: Thu, 24 Oct 2024 20:41:44 -0700
Subject: [PATCH] Extend the compact gate implementation to provide `index`
 method

This method essentially returns the unique index associated with the current gate. The plan is to use that number for metrics uniqueness check: they greatly benefit from knowing that the dimension for compact gate is `u64`
---
 ipa-step-derive/src/lib.rs | 12 ++++++++++++
 ipa-step-test/src/lib.rs   | 13 +++++++++++++
 2 files changed, 25 insertions(+)

diff --git a/ipa-step-derive/src/lib.rs b/ipa-step-derive/src/lib.rs
index ea5d19268..2298e3969 100644
--- a/ipa-step-derive/src/lib.rs
+++ b/ipa-step-derive/src/lib.rs
@@ -165,6 +165,18 @@ fn derive_gate_impl(ast: &DeriveInput) -> TokenStream {
                 <Self as ::std::fmt::Display>::fmt(self, f)
             }
         }
+
+        impl #name {
+            /// Returns the current index. It matches the index of the latest step
+            /// this gate has been narrowed to.
+            ///
+            /// If gate hasn't been narrowed yet, it returns the index of the default value.
+            #[must_use]
+            pub fn index(&self) -> ::ipa_step::CompactGateIndex {
+                self.0
+            }
+        }
+
     };
 
     // This environment variable is set by build scripts,
diff --git a/ipa-step-test/src/lib.rs b/ipa-step-test/src/lib.rs
index 84eab760d..3789e6fcf 100644
--- a/ipa-step-test/src/lib.rs
+++ b/ipa-step-test/src/lib.rs
@@ -17,15 +17,21 @@ mod tests {
 
     #[test]
     fn narrows() {
+        assert_eq!(ComplexGate::default().index(), 0);
         assert_eq!(ComplexGate::default().as_ref(), "/");
         assert_eq!(
             ComplexGate::default().narrow(&ComplexStep::One).as_ref(),
             "/one"
         );
+        assert_eq!(ComplexGate::default().narrow(&ComplexStep::One).index(), 1,);
         assert_eq!(
             ComplexGate::default().narrow(&ComplexStep::Two(2)).as_ref(),
             "/two2"
         );
+        assert_eq!(
+            ComplexGate::default().narrow(&ComplexStep::Two(2)).index(),
+            10,
+        );
         assert_eq!(
             ComplexGate::default()
                 .narrow(&ComplexStep::Two(2))
@@ -33,6 +39,13 @@ mod tests {
                 .as_ref(),
             "/two2/one"
         );
+        assert_eq!(
+            ComplexGate::default()
+                .narrow(&ComplexStep::Two(2))
+                .narrow(&BasicStep::One)
+                .index(),
+            11,
+        );
         assert_eq!(
             ComplexGate::from("/two2/one"),
             ComplexGate::default()