Include precedence calculation in ParseItemSet::transitions

spec/compiler/build_tables/parse_item_spec.cc

@@ -9,25 +9,30 @@ using namespace build_tables;
 START_TEST
 
 describe("ParseItemSet::transitions())", [&]() {
-  it("returns a map of symbols to ParseItemSets that would result from consuming those symbols", [&]() {
+  it("computes the ParseItemSet that would occur after consuming each lookahead symbol, along with its precedence", [&]() {
     SyntaxGrammar grammar{{
       SyntaxVariable("rule_0", VariableTypeNamed, {
         Production({
           {Symbol(11, true), 0, AssociativityNone, 101},
           {Symbol(12, true), 0, AssociativityNone, 102},
-          {Symbol(1), 0, AssociativityNone, 103},
-          {Symbol(13, true), 0, AssociativityNone, 104},
+          {Symbol(13), 5, AssociativityNone, 103},
+          {Symbol(14, true), 0, AssociativityNone, 104},
+        }),
+        Production({
+          {Symbol(11, true), 0, AssociativityNone, 105},
+          {Symbol(12, true), 0, AssociativityNone, 106},
+          {Symbol(15), 6, AssociativityNone, 107},
         })
       }),
       SyntaxVariable("rule_1", VariableTypeNamed, {
         Production({
-          {Symbol(2), 0, AssociativityNone, 105},
-          {Symbol(14, true), 0, AssociativityNone, 106},
+          {Symbol(15), 7, AssociativityNone, 109},
+          {Symbol(16, true), 0, AssociativityNone, 110},
         })
       }),
       SyntaxVariable("rule_2", VariableTypeNamed, {
         Production({
-          {Symbol(15, true), 0, AssociativityNone, 105},
+          {Symbol(18, true), 0, AssociativityNone, 111},
         })
       })
     }, {}, {}};
@@ -36,41 +41,69 @@ describe("ParseItemSet::transitions())", [&]() {
       return grammar.variables[variable_index].productions[production_index];
     };
 
-    ParseItemSet set1({
+    ParseItemSet item_set({
+
+      // Two symbols into the first production for rule_0
       {
         ParseItem(Symbol(0), production(0, 0), 2),
-        LookaheadSet({ Symbol(16, true) })
+        LookaheadSet({ Symbol(21, true) })
       },
+
+      // Two symbols into the second production for rule_0
+      {
+        ParseItem(Symbol(0), production(0, 1), 2),
+        LookaheadSet({ Symbol(21, true) })
+      },
+
+      // At the beginning of the first production for rule_1
       {
         ParseItem(Symbol(1), production(1, 0), 0),
-        LookaheadSet({ Symbol(17, true) })
+        LookaheadSet({ Symbol(22, true) })
       },
+
+      // At the end of the first production for rule_2
       {
         ParseItem(Symbol(2), production(2, 0), 1),
-        LookaheadSet({ Symbol(17, true) })
+        LookaheadSet({ Symbol(22, true) })
       }
     });
 
-    AssertThat(set1.transitions(), Equals(map<Symbol, ParseItemSet>({
+    AssertThat(item_set.transitions(), Equals(ParseItemSet::TransitionMap({
+
+      // For the first item, symbol 13 is next, with precedence 5.
       {
-        Symbol(1),
-        ParseItemSet({
-          {
-            ParseItem(Symbol(0), production(0, 0), 3),
-            LookaheadSet({ Symbol(16, true) })
-          }
-        })
+        Symbol(13),
+        {
+          ParseItemSet({
+            {
+              ParseItem(Symbol(0), production(0, 0), 3),
+              LookaheadSet({ Symbol(21, true) })
+            }
+          }),
+          PrecedenceRange(5, 5)
+        }
       },
 
+      // For the second and third item, symbol 15 is next, with two different
+      // precedence values.
       {
-        Symbol(2),
-        ParseItemSet({
-          {
-            ParseItem(Symbol(1), production(1, 0), 1),
-            LookaheadSet({ Symbol(17, true) })
-          },
-        })
+        Symbol(15),
+        {
+          ParseItemSet({
+            {
+              ParseItem(Symbol(0), production(0, 1), 3),
+              LookaheadSet({ Symbol(21, true) })
+            },
+            {
+              ParseItem(Symbol(1), production(1, 0), 1),
+              LookaheadSet({ Symbol(22, true) })
+            },
+          }),
+          PrecedenceRange(6, 7)
+        }
       },
+
+      // The third item is at the end of its production: no transitions.
     })));
   });
 });

src/compiler/build_tables/build_parse_table.cc

@@ -99,11 +99,12 @@ class ParseTableBuilder {
   void add_shift_actions(const ParseItemSet &item_set, ParseStateId state_id) {
     for (const auto &transition : item_set.transitions()) {
       const Symbol &symbol = transition.first;
-      const ParseItemSet &next_item_set = transition.second;
+      const ParseItemSet &next_item_set = transition.second.first;
+      const PrecedenceRange &precedence = transition.second.second;
 
       ParseAction *new_action = add_action(
         state_id, symbol,
-        ParseAction::Shift(0, precedence_values_for_item_set(next_item_set)),
+        ParseAction::Shift(0, precedence),
         item_set);
       if (new_action)
         new_action->state_index = add_parse_state(next_item_set);
@@ -319,16 +320,6 @@ class ParseTableBuilder {
     return result;
   }
 
-  PrecedenceRange precedence_values_for_item_set(const ParseItemSet &item_set) {
-    PrecedenceRange result;
-    for (const auto &pair : item_set.entries) {
-      const ParseItem &item = pair.first;
-      if (item.step_index > 0)
-        result.add(item.production->at(item.step_index - 1).precedence);
-    }
-    return result;
-  }
-
   string symbol_name(const rules::Symbol &symbol) const {
     if (symbol.is_built_in()) {
       if (symbol == rules::ERROR())

src/compiler/build_tables/parse_item.cc

@@ -104,8 +104,8 @@ size_t ParseItemSet::Hash::operator()(const ParseItemSet &item_set) const {
   return result;
 }
 
-map<Symbol, ParseItemSet> ParseItemSet::transitions() const {
-  map<Symbol, ParseItemSet> result;
+ParseItemSet::TransitionMap ParseItemSet::transitions() const {
+  ParseItemSet::TransitionMap result;
   for (const auto &pair : entries) {
     const ParseItem &item = pair.first;
     const LookaheadSet &lookahead_symbols = pair.second;
@@ -114,9 +114,11 @@ map<Symbol, ParseItemSet> ParseItemSet::transitions() const {
 
     size_t step = item.step_index + 1;
     Symbol symbol = item.production->at(item.step_index).symbol;
+    int precedence = item.production->at(item.step_index).precedence;
     ParseItem new_item(item.lhs(), *item.production, step);
 
-    result[symbol].entries[new_item] = lookahead_symbols;
+    result[symbol].first.entries[new_item] = lookahead_symbols;
+    result[symbol].second.add(precedence);
   }
 
   return result;

src/compiler/build_tables/parse_item.h

@@ -7,6 +7,7 @@
 #include "compiler/rules/symbol.h"
 #include "compiler/rules/metadata.h"
 #include "compiler/syntax_grammar.h"
+#include "compiler/precedence_range.h"
 
 namespace tree_sitter {
 namespace build_tables {
@@ -37,7 +38,9 @@ class ParseItemSet {
   ParseItemSet();
   explicit ParseItemSet(const std::map<ParseItem, LookaheadSet> &);
 
-  std::map<rules::Symbol, ParseItemSet> transitions() const;
+  typedef std::map<rules::Symbol, std::pair<ParseItemSet, PrecedenceRange>> TransitionMap;
+
+  TransitionMap transitions() const;
   bool operator==(const ParseItemSet &) const;
 
   std::map<ParseItem, LookaheadSet> entries;