From 8d70ef3a80bfe34c12ccf96c2b43717c4e056255 Mon Sep 17 00:00:00 2001
From: "Jip J. Dekker" <jip@dekker.one>
Date: Thu, 15 Feb 2024 11:20:37 +0100
Subject: [PATCH] Add index and find operations for VarRange

---
 crates/pindakaas/src/helpers.rs | 28 +++++++++++++++++++++++++++-
 1 file changed, 27 insertions(+), 1 deletion(-)

diff --git a/crates/pindakaas/src/helpers.rs b/crates/pindakaas/src/helpers.rs
index f2dbef6c63..cbf3c31772 100644
--- a/crates/pindakaas/src/helpers.rs
+++ b/crates/pindakaas/src/helpers.rs
@@ -2,6 +2,7 @@ use std::{
 	cmp::max,
 	collections::HashSet,
 	iter::FusedIterator,
+	num::NonZeroI32,
 	ops::{Bound, RangeBounds, RangeInclusive},
 };
 
@@ -208,9 +209,34 @@ pub struct VarRange {
 }
 
 impl VarRange {
+	/// Create a range starting from [`start`] and ending at [`end`] (inclusive)
 	pub fn new(start: Var, end: Var) -> Self {
 		Self { start, end }
 	}
+
+	/// Performs the indexing operation into the variable range
+	pub fn index(&self, index: usize) -> Var {
+		if index >= self.len() {
+			panic!("out of bounds access");
+		}
+		if index == 0 {
+			self.start
+		} else {
+			let index = NonZeroI32::new(index as i32).unwrap();
+			self.start.checked_add(index).unwrap()
+		}
+	}
+
+	/// Find the index of a variable within the range
+	pub fn find(&self, var: Var) -> Option<usize> {
+		if !self.contains(&var) {
+			None
+		} else {
+			let offset = (var.0.get() - self.start.0.get()) as usize;
+			debug_assert!(offset <= self.len());
+			Some(offset)
+		}
+	}
 }
 
 impl Iterator for VarRange {
@@ -226,7 +252,7 @@ impl Iterator for VarRange {
 		}
 	}
 	fn size_hint(&self) -> (usize, Option<usize>) {
-		let size = max(self.end.0.get() - self.start.0.get(), 0) as usize;
+		let size = max(self.end.0.get() - self.start.0.get() + 1, 0) as usize;
 		(size, Some(size))
 	}
 	fn count(self) -> usize {