From 4563cbe7130158952df99b3517a5671c19ca7187 Mon Sep 17 00:00:00 2001
From: Jake Hillion <jakehillion@meta.com>
Date: Thu, 12 Oct 2023 21:33:39 -0700
Subject: [PATCH] tbv2: improve equality for iterator

---
 include/oi/IntrospectionResult-inl.h | 14 +++++++++++---
 include/oi/IntrospectionResult.h     |  6 ++++++
 oi/IntrospectionResult.cpp           |  6 +++++-
 3 files changed, 22 insertions(+), 4 deletions(-)

diff --git a/include/oi/IntrospectionResult-inl.h b/include/oi/IntrospectionResult-inl.h
index 08cd451..13928b5 100644
--- a/include/oi/IntrospectionResult-inl.h
+++ b/include/oi/IntrospectionResult-inl.h
@@ -72,9 +72,17 @@ inline const result::Element* IntrospectionResult::const_iterator::operator->()
 
 inline bool IntrospectionResult::const_iterator::operator==(
     const IntrospectionResult::const_iterator& that) const {
-  return this->data_ == that.data_ && !this->next_.has_value() &&
-         !that.next_.has_value();  // TODO: is this sufficient? kind of hacky as
-                                   // this only works for comparing to .end()
+  // Case 1: The next data to read differs, thus the iterators are different.
+  if (this->data_ != that.data_)
+    return false;
+  // Case 2: Both iterators have no next value, thus they are both complete. It
+  // is insufficient to check increments as the number of increments is unknown
+  // when constructing `end()`.
+  if (!this->next_.has_value() && !that.next_.has_value())
+    return true;
+  // Case 3: The iterators are reading the same data. If they have produced the
+  // same number of elements they are equal, else they are not.
+  return this->increments_ == that.increments_;
 }
 inline bool IntrospectionResult::const_iterator::operator!=(
     const IntrospectionResult::const_iterator& that) const {
diff --git a/include/oi/IntrospectionResult.h b/include/oi/IntrospectionResult.h
index 8674649..427463b 100644
--- a/include/oi/IntrospectionResult.h
+++ b/include/oi/IntrospectionResult.h
@@ -63,6 +63,12 @@ class IntrospectionResult {
     // improvement but it isn't copyable. A string type with size fixed at
     // construction would also be good.
     std::list<std::string> dynamic_type_path_;
+
+    // We cannot track the position in the iteration solely by the underlying
+    // iterator as some fields do not extract data (for example, primitives).
+    // Track the number of increment operations as well to get an accurate
+    // equality check.
+    uint64_t increments_ = 0;
   };
 
   IntrospectionResult(std::vector<uint8_t> buf, exporters::inst::Inst inst);
diff --git a/oi/IntrospectionResult.cpp b/oi/IntrospectionResult.cpp
index bc6cf22..efb11a2 100644
--- a/oi/IntrospectionResult.cpp
+++ b/oi/IntrospectionResult.cpp
@@ -30,9 +30,13 @@ namespace oi {
 IntrospectionResult::const_iterator&
 IntrospectionResult::const_iterator::operator++() {
   if (stack_.empty()) {
-    next_ = std::nullopt;
+    if (next_ != std::nullopt) {
+      ++increments_;
+      next_ = std::nullopt;
+    }
     return *this;
   }
+  ++increments_;
 
   auto el = stack_.top();
   stack_.pop();