From e79d8872acd76b70e7d419fd06f68b6cf3996419 Mon Sep 17 00:00:00 2001
From: Michael Miller <icy.arctic.fox@gmail.com>
Date: Mon, 19 Nov 2018 19:41:51 -0700
Subject: [PATCH] Cleanup indexer code

---
 src/spectator/example_group.cr | 37 ++++++++++++++++++++--------------
 1 file changed, 22 insertions(+), 15 deletions(-)

diff --git a/src/spectator/example_group.cr b/src/spectator/example_group.cr
index a1e2e2e..ee3494b 100644
--- a/src/spectator/example_group.cr
+++ b/src/spectator/example_group.cr
@@ -31,25 +31,32 @@ module Spectator
     getter example_count = 0
 
     def [](index : Int) : Example
-      offset = index
-      offset += example_count if offset < 0
-      raise IndexError.new if offset < 0 || offset >= example_count
-      found = children.find do |child|
-        count = child.example_count
-        if offset < count
-          true
-        else
-          offset -= count
-          false
-        end
-      end
-      if found
-        found[offset]
+      offset = check_bounds(index)
+      find_nested(offset)
+    end
+
+    private def check_bounds(index)
+      if index < 0
+        raise IndexError.new if index < -example_count
+        index + example_count
       else
-        raise IndexError.new
+        raise IndexError.new if index >= example_count
+        index
       end
     end
 
+    private def find_nested(index)
+      offset = index
+      child = children.each do |child|
+        count = child.example_count
+        break child if offset < count
+        offset -= count
+      end
+      # It should be impossible to get `nil` here,
+      # provided the bounds check and example counts are correct.
+      child.not_nil![offset]
+    end
+
     def finished? : Bool
       children.all?(&.finished?)
     end