From a2795bb5312ead662a363b8e6358c83906ebb15a Mon Sep 17 00:00:00 2001
From: Josh Holtrop <jholtrop@gmail.com>
Date: Tue, 21 Sep 2021 17:09:53 -0400
Subject: [PATCH] Keep track of follow item sets by symbol for each item set

---
 lib/imbecile/parser.rb          | 43 +++++++++++++++++++++++++--------
 lib/imbecile/parser/item_set.rb |  7 +++++-
 2 files changed, 39 insertions(+), 11 deletions(-)

diff --git a/lib/imbecile/parser.rb b/lib/imbecile/parser.rb
index 6bb8de9..9c43bf2 100644
--- a/lib/imbecile/parser.rb
+++ b/lib/imbecile/parser.rb
@@ -3,11 +3,11 @@ class Imbecile
   class Parser
 
     def initialize(tokens, rules)
-      token_eof = Token.new("$", nil, TOKEN_EOF)
+      @token_eof = Token.new("$", nil, TOKEN_EOF)
       @item_sets = []
-      item_sets_set = Set.new
+      @item_sets_set = {}
       start_items = rules["Start"].map do |rule|
-        rule.components << token_eof
+        rule.components << @token_eof
         Item.new(rule, 0)
       end
       eval_item_sets = Set.new
@@ -17,22 +17,45 @@ class Imbecile
         this_eval_item_sets = eval_item_sets
         eval_item_sets = Set.new
         this_eval_item_sets.each do |item_set|
-          unless item_sets_set.include?(item_set)
+          unless @item_sets_set.include?(item_set)
             item_set.id = @item_sets.size
             @item_sets << item_set
-            item_sets_set << item_set
-            puts "Item set #{item_set.id}:"
-            puts item_set
-            puts
+            @item_sets_set[item_set] = item_set
             item_set.follow_symbols.each do |follow_symbol|
-              unless follow_symbol == token_eof
-                follow_set = item_set.follow_set(follow_symbol)
+              unless follow_symbol == @token_eof
+                follow_set = item_set.build_follow_set(follow_symbol)
                 eval_item_sets << follow_set
               end
             end
           end
         end
       end
+
+      @item_sets.each do |item_set|
+        process_item_set(item_set)
+        puts "Item set #{item_set.id}:"
+        puts item_set
+        item_set.follow_item_set.each do |follow_symbol, follow_item_set|
+          if follow_symbol.is_a?(Token)
+            name = follow_symbol.name
+          else
+            name = follow_symbol[0].name
+          end
+          puts " #{name} => #{follow_item_set.id}"
+        end
+        puts
+      end
+    end
+
+    private
+
+    def process_item_set(item_set)
+      item_set.follow_symbols.each do |follow_symbol|
+        unless follow_symbol == @token_eof
+          follow_set = @item_sets_set[item_set.build_follow_set(follow_symbol)]
+          item_set.follow_item_set[follow_symbol] = follow_set
+        end
+      end
     end
 
   end
diff --git a/lib/imbecile/parser/item_set.rb b/lib/imbecile/parser/item_set.rb
index 619c222..86d683a 100644
--- a/lib/imbecile/parser/item_set.rb
+++ b/lib/imbecile/parser/item_set.rb
@@ -7,8 +7,13 @@ class Imbecile
 
       attr_accessor :id
 
+      # @return [Hash]
+      #   Maps a follow symbol to its item set.
+      attr_reader :follow_item_set
+
       def initialize(items)
         @items = Set.new(items)
+        @follow_item_set = {}
         close!
       end
 
@@ -16,7 +21,7 @@ class Imbecile
         Set.new(@items.map(&:follow_symbol).compact)
       end
 
-      def follow_set(symbol)
+      def build_follow_set(symbol)
         ItemSet.new(items_followed_by(symbol).map(&:next_position))
       end