[pseudo] Check follow-sets instead of tying reduce actions to lookahead tokens.

[sammccall]

Summary

Previously, the action table stores a reduce action for each lookahead
token it should allow. These tokens are the followSet(action.rule.target).

In practice, the follow sets are large, so we spend a bunch of time binary
searching around all these essentially-duplicates to check whether our lookahead
token is there.
However the number of reduces for a given state is very small, so we're
much better off linear scanning over them and performing a fast check for each.

D128318 was an attempt at this, storing a bitmap for each reduce.
However it's even more compact just to use the follow sets directly, as
there are fewer nonterminals than (state, rule) pairs. It's also faster.

This specialized approach means unbundling Reduce from other actions in
LRTable, so it's no longer useful to support it in Action. I suspect
Action will soon go away, as we store each kind of action separately.

This improves glrParse speed by 42% (3.30 -> 4.69 MB/s).
It also reduces LR table size by 59% (343 -> 142kB).

[hokein]

This looks a great improvement. A few nits, it looks good to me overall.

clang-tools-extra/pseudo/include/clang-pseudo/grammar/LRTable.h

136

I think the API is motivated by the LR(0) experimentation, we're now decoupling the reduction look up from the lookahead token, this seems nice (LR(0) only needs getReduceRules, other LR(1) variant need getReduceRules + canFollow).

The API usage for SLR(1) is not quite straight forward at the first glance, would be better to add a small snippet in the comment.

143

nit: use isNonterminal(Nonterminal)? IMO it is clearer.

215

nit: I think rephasing it as: a half-open range of Reduces [ReduceOffset[S], ReduceOffset[S+1]). is clearer -- the current syntax Reduces[ [ReduceOffset[S], ReduceOffset[S+1]) ] is a bit weird (My first attempt was to read it as an unmatch-brackets code snippet)

218

The NT is not quite obvious, what is NT? I think it is the symbol ID, then I don't understand the NT * NUM_TOKENS + tok::Kind(T) part.

219

Oh, I understand it now after reading the build code...

NT here is the non-terminal (symbol to be reduced), T is the lookahead. And FollowSets is conceptually a flat array of a 2d bool array [Nonterminal, Tokens].

Probably add some more comments.

clang-tools-extra/pseudo/lib/GLR.cpp

371

nit: add a comment about what does return value indicate.

clang-tools-extra/pseudo/lib/grammar/LRTableBuild.cpp

115

make these members as private -- having two public members seems to violate the pattern. (FollowSets can be initialized in the constructor, we might need a separate method for Reduces. or either change the Builder to a struct).

clang-tools-extra/pseudo/unittests/LRTableTest.cpp

28

nit: we use the hard-coded ruleID in the following test, I'd suggest add a trailing comment for the rule ID for each rule here, for better code readability.

38

It would be better to change it to an IDENTIFIER to reflect the grammar above.

39

nit: I'd move this on Line67, as this token serves as an invalid token according to the grammar.

42

the row is usual for terminals, and the reduce is a separate table, which seems a bit confusing -- I would probably split the reduce table out (adding an empty col between term and reduce)

45

nit: acc => R2 (acc) (sorry, I didn't update the comment in the removing-accept patch)

52

nit: add a /*RuleID=*/ comment annotation for the second argument.

[sammccall]

clang-tools-extra/pseudo/include/clang-pseudo/grammar/LRTable.h

136

Added.

In fact it wasn't motivated by the LR(0) stuff, at least not consciously/directly.
Rather I just wanted to avoid having to allocate a data structure to copy a filtered list into!

clang-tools-extra/pseudo/lib/grammar/LRTableBuild.cpp

115

This whole class is private and exists purely for the convenience of buildSLR()/buildForTests, I don't think adding wrapping methods really helps them out.

Changed Builder to a struct and made everything public.

clang-tools-extra/pseudo/unittests/LRTableTest.cpp

42

I moved it to a separate (1D) table above ReduceEntries, and at that point it looked exactly like the definition of ReduceEntries.

So I don't think it adds anything that's not obvious from the code, I dropped the comment.

[sammccall]

address comments

[bjope]

Our downstream bots have failed when using LLVM_ENABLE_EXPENSIVE_CHECKS=ON , and I think (unless my bisecting got screwed up) that it started to happen after this patch.

We typically get

Failed Tests (10):
  ClangPseudo :: lr-build-basic.test
  ClangPseudo :: lr-build-conflicts.test
  clangPseudo Unit Tests :: ./ClangPseudoTests/10/31
  clangPseudo Unit Tests :: ./ClangPseudoTests/11/31
  clangPseudo Unit Tests :: ./ClangPseudoTests/12/31
  clangPseudo Unit Tests :: ./ClangPseudoTests/13/31
  clangPseudo Unit Tests :: ./ClangPseudoTests/14/31
  clangPseudo Unit Tests :: ./ClangPseudoTests/15/31
  clangPseudo Unit Tests :: ./ClangPseudoTests/16/31
  clangPseudo Unit Tests :: ./ClangPseudoTests/9/31

And the failing stack traces look like this:

/lib/gcc/x86_64-pc-linux-gnu/9.3.0/../../../../include/c++/9.3.0/bits/stl_vector.h:1060: std::vector::const_reference std::vector<unsigned short, std::allocator<unsigned short> >::operator[](std::vector::size_type) const [_Tp = unsigned short, _Alloc = std::allocator<unsigned short>]: Assertion '__builtin_expect(__n < this->size(), true)' failed.
 #0 0x000000000047b6b3 llvm::sys::PrintStackTrace(llvm::raw_ostream&, int) (/repo/llvm/build-all-expensive/tools/clang/tools/extra/pseudo/unittests/./ClangPseudoTests+0x47b6b3)
 #1 0x00000000004795cc llvm::sys::RunSignalHandlers() (/repo/llvm/build-all-expensive/tools/clang/tools/extra/pseudo/unittests/./ClangPseudoTests+0x4795cc)
 #2 0x000000000047bcc6 SignalHandler(int) (/repo/llvm/build-all-expensive/tools/clang/tools/extra/pseudo/unittests/./ClangPseudoTests+0x47bcc6)
 #3 0x00007f9ace5af630 __restore_rt (/lib64/libpthread.so.0+0xf630)
 #4 0x00007f9acd6ce387 __GI_raise (/lib64/libc.so.6+0x36387)
 #5 0x00007f9acd6cfa78 __GI_abort (/lib64/libc.so.6+0x37a78)
 #6 0x000000000045b988 (/repo/llvm/build-all-expensive/tools/clang/tools/extra/pseudo/unittests/./ClangPseudoTests+0x45b988)
 #7 0x00000000004e42a6 clang::pseudo::(anonymous namespace)::GLRReduce::popPending() (/repo/llvm/build-all-expensive/tools/clang/tools/extra/pseudo/unittests/./ClangPseudoTests+0x4e42a6)
 #8 0x00000000004e15f2 clang::pseudo::(anonymous namespace)::GLRReduce::operator()(std::vector<clang::pseudo::GSS::Node const*, std::allocator<clang::pseudo::GSS::Node const*> >&, unsigned short) (/repo/llvm/build-all-expensive/tools/clang/tools/extra/pseudo/unittests/./ClangPseudoTests+0x4e15f2)
 #9 0x00000000004e2e15 clang::pseudo::glrReduce(std::vector<clang::pseudo::GSS::Node const*, std::allocator<clang::pseudo::GSS::Node const*> >&, unsigned short, clang::pseudo::ParseParams const&) (/repo/llvm/build-all-expensive/tools/clang/tools/extra/pseudo/unittests/./ClangPseudoTests+0x4e2e15)
#10 0x0000000000432379 clang::pseudo::(anonymous namespace)::GLRTest_ReduceConflictsSplitting_Test::TestBody() (/repo/llvm/build-all-expensive/tools/clang/tools/extra/pseudo/unittests/./ClangPseudoTests+0x432379)
#11 0x00000000004b8b6c testing::Test::Run() (/repo/llvm/build-all-expensive/tools/clang/tools/extra/pseudo/unittests/./ClangPseudoTests+0x4b8b6c)

[hokein]

Thanks for the report. There is an out-of-bound issue in LRTable::getReduceRules, fixed in c99827349927a44334f2b04139168efd0bc87cd3.

[bjope]

In D128472#3625584, @hokein wrote:

Thanks for the report. There is an out-of-bound issue in LRTable::getReduceRules, fixed in c99827349927a44334f2b04139168efd0bc87cd3.

Great, thanks for the fix!

clang-tools-extra/pseudo/include/clang-pseudo/grammar/LRTable.h

 //
 //===----------------------------------------------------------------------===//
 
 #ifndef CLANG_PSEUDO_GRAMMAR_LRTABLE_H
 #define CLANG_PSEUDO_GRAMMAR_LRTABLE_H
 
 #include "clang-pseudo/grammar/Grammar.h"
 #include "llvm/ADT/ArrayRef.h"
+#include "llvm/ADT/BitVector.h"
+#include "llvm/Support/Capacity.h"
 #include <cstdint>
 #include <vector>
 
 namespace clang {
 namespace pseudo {
 
 // Represents the LR parsing table, which can efficiently the question "what is
 // the next step given the lookahead token and current state on top of the
 // stack?".
 //
 // This is a dense implementation, which only takes an amount of space that is
 // proportional to the number of non-empty entries in the table.
 //
 // Unlike the typical LR parsing table which allows at most one available action
 // per entry, conflicted actions are allowed in LRTable. The LRTable is designed
 // to be used in nondeterministic LR parsers (e.g. GLR).
 class LRTable {
 public:
   // StateID is only 13 bits wide.
   using StateID = uint16_t;
   static constexpr unsigned StateBits = 13;
 
   // Action represents the terminal and nonterminal actions, it combines the
   // entry of the ACTION and GOTO tables from the LR literature.
+  //
+  // FIXME: as we move away from a homogeneous table structure shared between
+  // action types, this class becomes less useful. Remove it.
   class Action {
   public:
     enum Kind : uint8_t {
       Sentinel = 0,
       // Terminal actions, corresponding to entries of ACTION table.
 
       // Shift to state n: move forward with the lookahead, and push state n
       // onto the state stack.
       // A shift is a forward transition, and the value n is the next state that
       // the parser is to enter.
       Shift,
-      // Reduce by a rule: pop the state stack.
-      Reduce,
 
       // NOTE: there are no typical accept actions in the LRtable, accept
       // actions are handled specifically in the parser -- if the parser
       // reaches to a target state (which is goto(StartState, StartSymbol)) at
       // the EOF token after a reduce, this indicates the input has been parsed
       // as the StartSymbol successfully.
 
       // Nonterminal actions, corresponding to entry of GOTO table.
 
       // Go to state n: push state n onto the state stack.
       // Similar to Shift, but it is a nonterminal forward transition.
       GoTo,
     };
 
     static Action goTo(StateID S) { return Action(GoTo, S); }
     static Action shift(StateID S) { return Action(Shift, S); }
-    static Action reduce(RuleID RID) { return Action(Reduce, RID); }
     static Action sentinel() { return Action(Sentinel, 0); }
 
     StateID getShiftState() const {
       assert(kind() == Shift);
       return Value;
     }
     StateID getGoToState() const {
       assert(kind() == GoTo);
       return Value;
     }
-    RuleID getReduceRule() const {
-      assert(kind() == Reduce);
-      return Value;
-    }
     Kind kind() const { return static_cast<Kind>(K); }
 
     bool operator==(const Action &L) const { return opaque() == L.opaque(); }
     uint16_t opaque() const { return K << ValueBits | Value; };
 
   private:
     Action(Kind K1, unsigned Value) : K(K1), Value(Value) {}
     static constexpr unsigned ValueBits = StateBits;
     static constexpr unsigned KindBits = 3;
     static_assert(ValueBits >= RuleBits, "Value must be able to store RuleID");
     static_assert(KindBits + ValueBits <= 16,
                   "Must be able to store kind and value efficiently");
     uint16_t K : KindBits;
     // Either StateID or RuleID, depending on the Kind.
     uint16_t Value : ValueBits;
   };
 
-  // Returns all available actions for the given state on a terminal.
-  // Expected to be called by LR parsers.
-  llvm::ArrayRef<Action> getActions(StateID State, SymbolID Terminal) const;
   // Returns the state after we reduce a nonterminal.
   // Expected to be called by LR parsers.
   // REQUIRES: Nonterminal is valid here.
   StateID getGoToState(StateID State, SymbolID Nonterminal) const;
   // Returns the state after we shift a terminal.
   // Expected to be called by LR parsers.
   // If the terminal is invalid here, returns None.
   llvm::Optional<StateID> getShiftState(StateID State, SymbolID Terminal) const;
 
-  // Looks up available actions.
-  // Returns empty if no available actions in the table.
-  llvm::ArrayRef<Action> find(StateID State, SymbolID Symbol) const;
+  // Returns the possible reductions from a state.
+  //
+  // These are not keyed by a lookahead token. Instead, call canFollow() to
+  // check whether a reduction should apply in the current context:
+  //   for (RuleID R : LR.getReduceRules(S)) {
+  //     if (!LR.canFollow(G.lookupRule(R).Target, NextToken))
+  //       continue;
+  //     // ...apply reduce...
+  //   }
+  llvm::ArrayRef<RuleID> getReduceRules(StateID State) const {
+    return llvm::makeArrayRef(&Reduces[ReduceOffset[State]],
+                              &Reduces[ReduceOffset[State + 1]]);
+  }
+  // Returns whether Terminal can follow Nonterminal in a valid source file.
+  bool canFollow(SymbolID Nonterminal, SymbolID Terminal) const {
+    assert(isToken(Terminal));
+    assert(isNonterminal(Nonterminal));
+    return FollowSets.test(tok::NUM_TOKENS * Nonterminal +
+                           symbolToToken(Terminal));
+  }
 
   // Returns the state from which the LR parser should start to parse the input
   // tokens as the given StartSymbol.
   //
   // In LR parsing, the start state of `translation-unit` corresponds to
   // `_ := • translation-unit`.
   //
   // Each start state responds to **a** single grammar rule like `_ := start`.
   // REQUIRE: The given StartSymbol must exist in the grammar (in a form of
   //          `_ := start`).
   StateID getStartState(SymbolID StartSymbol) const;
 
   size_t bytes() const {
-    return sizeof(*this) + Actions.capacity() * sizeof(Action) +
-           Symbols.capacity() * sizeof(SymbolID) +
-           StateOffset.capacity() * sizeof(uint32_t);
+    return sizeof(*this) + llvm::capacity_in_bytes(Actions) +
+           llvm::capacity_in_bytes(Symbols) +
+           llvm::capacity_in_bytes(StateOffset) +
+           llvm::capacity_in_bytes(Reduces) +
+           llvm::capacity_in_bytes(ReduceOffset) +
+           llvm::capacity_in_bytes(FollowSets);
   }
 
   std::string dumpStatistics() const;
   std::string dumpForTests(const Grammar &G) const;
 
   // Build a SLR(1) parsing table.
   static LRTable buildSLR(const Grammar &G);
 
-  class Builder;
+  struct Builder;
   // Represents an entry in the table, used for building the LRTable.
   struct Entry {
     StateID State;
     SymbolID Symbol;
     Action Act;
   };
+  struct ReduceEntry {
+    StateID State;
+    RuleID Rule;
+  };
   // Build a specifid table for testing purposes.
-  static LRTable buildForTests(const GrammarTable &, llvm::ArrayRef<Entry>);
+  static LRTable buildForTests(const Grammar &G, llvm::ArrayRef<Entry>,
+                               llvm::ArrayRef<ReduceEntry>);
 
 private:
+  // Looks up actions stored in the generic table.
+  llvm::ArrayRef<Action> find(StateID State, SymbolID Symbol) const;
+
   // Conceptually the LR table is a multimap from (State, SymbolID) => Action.
   // Our physical representation is quite different for compactness.
 
   // Index is StateID, value is the offset into Symbols/Actions
   // where the entries for this state begin.
   // Give a state id, the corresponding half-open range of Symbols/Actions is
   // [StateOffset[id], StateOffset[id+1]).
   std::vector<uint32_t> StateOffset;
   // Parallel to Actions, the value is SymbolID (columns of the matrix).
   // Grouped by the StateID, and only subranges are sorted.
   std::vector<SymbolID> Symbols;
   // A flat list of available actions, sorted by (State, SymbolID).
   std::vector<Action> Actions;
   // A sorted table, storing the start state for each target parsing symbol.
   std::vector<std::pair<SymbolID, StateID>> StartStates;
+
+  // Given a state ID S, the half-open range of Reduces is
+  // [ReduceOffset[S], ReduceOffset[S+1])
+  std::vector<uint32_t> ReduceOffset;
+  std::vector<RuleID> Reduces;
+  // Conceptually this is a bool[SymbolID][Token], each entry describing whether
+  // the grammar allows the (nonterminal) symbol to be followed by the token.
+  //
+  // This is flattened by encoding the (SymbolID Nonterminal, tok::Kind Token)
+  // as an index: Nonterminal * NUM_TOKENS + Token.
+  llvm::BitVector FollowSets;
 };
 llvm::raw_ostream &operator<<(llvm::raw_ostream &, const LRTable::Action &);
 
 } // namespace pseudo
 } // namespace clang
 
 #endif // CLANG_PSEUDO_GRAMMAR_LRTABLE_H

clang-tools-extra/pseudo/lib/GLR.cpp

while (!Sequences.empty()) {
       pushNext();
       popPending();
     }
   }
 
 private:
   // pop walks up the parent chain(s) for a reduction from Head by to Rule.
   // Once we reach the end, record the bases and sequences.
-  void pop(const GSS::Node *Head, RuleID RID) {
+  void pop(const GSS::Node *Head, RuleID RID, const Rule &Rule) {
     LLVM_DEBUG(llvm::dbgs() << "  Pop " << Params.G.dumpRule(RID) << "\n");
-    const auto &Rule = Params.G.lookupRule(RID);
     Family F{/*Start=*/0, /*Symbol=*/Rule.Target, /*Rule=*/RID};
     TempSequence.resize_for_overwrite(Rule.Size);
     auto DFS = [&](const GSS::Node *N, unsigned I, auto &DFS) {
       TempSequence[Rule.Size - 1 - I] = N->Payload;
       if (I + 1 == Rule.Size) {
         F.Start = TempSequence.front()->startTokenIndex();
         LLVM_DEBUG({
           for (const auto *B : N->parents())
private:
   }
 
   // popPending pops every available reduction.
   void popPending() {
     for (; NextPopHead < Heads->size(); ++NextPopHead) {
       // In trivial cases, we perform the complete reduce here!
       if (popAndPushTrivial())
         continue;
-      for (const auto &A :
-           Params.Table.getActions((*Heads)[NextPopHead]->State, Lookahead)) {
-        if (A.kind() != LRTable::Action::Reduce)
-          continue;
-        pop((*Heads)[NextPopHead], A.getReduceRule());
+      for (RuleID RID :
+           Params.Table.getReduceRules((*Heads)[NextPopHead]->State)) {
+        const auto &Rule = Params.G.lookupRule(RID);
+        if (Params.Table.canFollow(Rule.Target, Lookahead))
+          pop((*Heads)[NextPopHead], RID, Rule);
       }
     }
   }
 
   // Storage reused by each call to pushNext.
   std::vector<std::pair</*Goto*/ StateID, const GSS::Node *>> FamilyBases;
   std::vector<std::pair<RuleID, SequenceRef>> FamilySequences;
   std::vector<const GSS::Node *> Parents;
private:
   // we pop and immediately push, as an LR parser (as opposed to GLR) would.
   // This is valid whenever there's only one concurrent PushSpec.
   //
   // This function handles a trivial but common subset of these cases:
   //  - there must be no pending pushes, and only one poppable head
   //  - the head must have only one reduction rule
   //  - the reduction path must be a straight line (no multiple parents)
   // (Roughly this means there's no local ambiguity, so the LR algorithm works).
+  //
+  // Returns true if we successfully consumed the next unpopped head.
   bool popAndPushTrivial() {
     if (!Sequences.empty() || Heads->size() != NextPopHead + 1)
       return false;
     const GSS::Node *Head = Heads->back();
     llvm::Optional<RuleID> RID;
-    for (auto &A : Params.Table.getActions(Head->State, Lookahead)) {
-      if (A.kind() != LRTable::Action::Reduce)
-        continue;
-      if (RID)
+    for (RuleID R : Params.Table.getReduceRules(Head->State)) {
+      if (RID.hasValue())
         return false;
-      RID = A.getReduceRule();
+      RID = R;
     }
     if (!RID)
       return true; // no reductions available, but we've processed the head!
     const auto &Rule = Params.G.lookupRule(*RID);
+    if (!Params.Table.canFollow(Rule.Target, Lookahead))
+      return true; // reduction is not available
     const GSS::Node *Base = Head;
     TempSequence.resize_for_overwrite(Rule.Size);
     for (unsigned I = 0; I < Rule.Size; ++I) {
       if (Base->parents().size() != 1)
         return false;
       TempSequence[Rule.Size - 1 - I] = Base->Payload;
       Base = Base->parents().front();
     }

clang-tools-extra/pseudo/lib/grammar/LRTable.cpp

 //===--- LRTable.cpp - Parsing table for LR parsers --------------*- C++-*-===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 
 #include "clang-pseudo/grammar/LRTable.h"
 #include "clang-pseudo/grammar/Grammar.h"
 #include "llvm/ADT/ArrayRef.h"
 #include "llvm/ADT/STLExtras.h"
+#include "llvm/ADT/StringExtras.h"
 #include "llvm/Support/ErrorHandling.h"
 #include "llvm/Support/FormatVariadic.h"
 #include "llvm/Support/raw_ostream.h"
 
 namespace clang {
 namespace pseudo {
 
 llvm::raw_ostream &operator<<(llvm::raw_ostream &OS, const LRTable::Action &A) {
   switch (A.kind()) {
   case LRTable::Action::Shift:
     return OS << llvm::formatv("shift state {0}", A.getShiftState());
-  case LRTable::Action::Reduce:
-    return OS << llvm::formatv("reduce by rule {0}", A.getReduceRule());
   case LRTable::Action::GoTo:
     return OS << llvm::formatv("go to state {0}", A.getGoToState());
   case LRTable::Action::Sentinel:
     llvm_unreachable("unexpected Sentinel action kind!");
   }
   llvm_unreachable("unexpected action kind!");
 }
 
 std::string LRTable::dumpStatistics() const {
   return llvm::formatv(R"(
 Statistics of the LR parsing table:
     number of states: {0}
     number of actions: {1}
-    size of the table (bytes): {2}
+    number of reduces: {2}
+    size of the table (bytes): {3}
 )",
-                       StateOffset.size() - 1, Actions.size(), bytes())
+                       StateOffset.size() - 1, Actions.size(), Reduces.size(),
+                       bytes())
       .str();
 }
 
 std::string LRTable::dumpForTests(const Grammar &G) const {
   std::string Result;
   llvm::raw_string_ostream OS(Result);
   OS << "LRTable:\n";
   for (StateID S = 0; S < StateOffset.size() - 1; ++S) {
     OS << llvm::formatv("State {0}\n", S);
     for (uint16_t Terminal = 0; Terminal < NumTerminals; ++Terminal) {
       SymbolID TokID = tokenSymbol(static_cast<tok::TokenKind>(Terminal));
       for (auto A : find(S, TokID)) {
         if (A.kind() == LRTable::Action::Shift)
-          OS.indent(4) << llvm::formatv("'{0}': shift state {1}\n",
+          OS.indent(4) << llvm::formatv("{0}: shift state {1}\n",
                                         G.symbolName(TokID), A.getShiftState());
-        else if (A.kind() == LRTable::Action::Reduce)
-          OS.indent(4) << llvm::formatv("'{0}': reduce by rule {1} '{2}'\n",
-                                        G.symbolName(TokID), A.getReduceRule(),
-                                        G.dumpRule(A.getReduceRule()));
       }
     }
+    for (RuleID R : getReduceRules(S)) {
+      SymbolID Target = G.lookupRule(R).Target;
+      std::vector<llvm::StringRef> Terminals;
+      for (unsigned Terminal = 0; Terminal < NumTerminals; ++Terminal) {
+        SymbolID TokID = tokenSymbol(static_cast<tok::TokenKind>(Terminal));
+        if (canFollow(Target, TokID))
+          Terminals.push_back(G.symbolName(TokID));
+      }
+      OS.indent(4) << llvm::formatv("{0}: reduce by rule {1} '{2}'\n",
+                                    llvm::join(Terminals, " "), R,
+                                    G.dumpRule(R));
+    }
     for (SymbolID NontermID = 0; NontermID < G.table().Nonterminals.size();
          ++NontermID) {
       if (find(S, NontermID).empty())
         continue;
-      OS.indent(4) << llvm::formatv("'{0}': go to state {1}\n",
+      OS.indent(4) << llvm::formatv("{0}: go to state {1}\n",
                                     G.symbolName(NontermID),
                                     getGoToState(S, NontermID));
     }
   }
   return OS.str();
 }
 
 llvm::Optional<LRTable::StateID>
 LRTable::getShiftState(StateID State, SymbolID Terminal) const {
   // FIXME: we spend a significant amount of time on misses here.
   // We could consider storing a std::bitset for a cheaper test?
   assert(pseudo::isToken(Terminal) && "expected terminal symbol!");
-  for (const auto &Result : getActions(State, Terminal))
+  for (const auto &Result : find(State, Terminal))
     if (Result.kind() == Action::Shift)
       return Result.getShiftState(); // unique: no shift/shift conflicts.
   return llvm::None;
 }
 
-llvm::ArrayRef<LRTable::Action> LRTable::getActions(StateID State,
-                                                    SymbolID Terminal) const {
-  assert(pseudo::isToken(Terminal) && "expect terminal symbol!");
-  return find(State, Terminal);
-}
-
 LRTable::StateID LRTable::getGoToState(StateID State,
                                        SymbolID Nonterminal) const {
   assert(pseudo::isNonterminal(Nonterminal) && "expected nonterminal symbol!");
   auto Result = find(State, Nonterminal);
   assert(Result.size() == 1 && Result.front().kind() == Action::GoTo);
   return Result.front().getGoToState();
 }
 

clang-tools-extra/pseudo/lib/grammar/LRTableBuild.cpp

 //===--- LRTableBuild.cpp - Build a LRTable from LRGraph ---------*- C++-*-===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 
 #include "clang-pseudo/grammar/Grammar.h"
 #include "clang-pseudo/grammar/LRGraph.h"
 #include "clang-pseudo/grammar/LRTable.h"
 #include "clang/Basic/TokenKinds.h"
+#include "llvm/ADT/SmallSet.h"
 #include <cstdint>
 
 namespace llvm {
 template <> struct DenseMapInfo<clang::pseudo::LRTable::Entry> {
   using Entry = clang::pseudo::LRTable::Entry;
   static inline Entry getEmptyKey() {
     static Entry E{static_cast<clang::pseudo::SymbolID>(-1), 0,
                    clang::pseudo::LRTable::Action::sentinel()};
return LHS.State == RHS.State && LHS.Symbol == RHS.Symbol &&
            LHS.Act == RHS.Act;
   }
 };
 } // namespace llvm
 
 namespace clang {
 namespace pseudo {
 
-class LRTable::Builder {
-public:
-  Builder(llvm::ArrayRef<std::pair<SymbolID, StateID>> StartStates)
-      : StartStates(StartStates) {}
+struct LRTable::Builder {
+  llvm::DenseMap<StateID, llvm::SmallSet<RuleID, 4>> Reduces;
+  std::vector<llvm::DenseSet<SymbolID>> FollowSets;
+  llvm::DenseSet<Entry> Entries;
+  std::vector<std::pair<SymbolID, StateID>> StartStates;
 
-  bool insert(Entry E) { return Entries.insert(std::move(E)).second; }
-  LRTable build(const GrammarTable &GT, unsigned NumStates) && {
+  LRTable build(unsigned NumStates) && {
     // E.g. given the following parsing table with 3 states and 3 terminals:
     //
     //            a    b     c
     // +-------+----+-------+-+
     // |state0 |    | s0,r0 | |
     // |state1 | acc|       | |
     // |state2 |    |  r1   | |
     // +-------+----+-------+-+
LRTable build(unsigned NumStates) && {
     Table.StateOffset = std::vector<uint32_t>(NumStates + 1, 0);
     size_t SortedIndex = 0;
     for (StateID State = 0; State < Table.StateOffset.size(); ++State) {
       Table.StateOffset[State] = SortedIndex;
       while (SortedIndex < Sorted.size() && Sorted[SortedIndex].State == State)
         ++SortedIndex;
     }
     Table.StartStates = std::move(StartStates);
-    return Table;
+
+    // Compile the follow sets into a bitmap.
+    Table.FollowSets.resize(tok::NUM_TOKENS * FollowSets.size());
+    for (SymbolID NT = 0; NT < FollowSets.size(); ++NT)
+      for (SymbolID Follow : FollowSets[NT])
+        Table.FollowSets.set(NT * tok::NUM_TOKENS + symbolToToken(Follow));
+
+    // Store the reduce actions in a vector partitioned by state.
+    Table.ReduceOffset.reserve(NumStates + 1);
+    std::vector<RuleID> StateRules;
+    for (StateID S = 0; S < NumStates; ++S) {
+      Table.ReduceOffset.push_back(Table.Reduces.size());
+      auto It = Reduces.find(S);
+      if (It == Reduces.end())
+        continue;
+      Table.Reduces.insert(Table.Reduces.end(), It->second.begin(),
+                           It->second.end());
+      std::sort(Table.Reduces.begin() + Table.ReduceOffset.back(),
+                Table.Reduces.end());
     }
+    Table.ReduceOffset.push_back(Table.Reduces.size());
 
-private:
-  llvm::DenseSet<Entry> Entries;
-  std::vector<std::pair<SymbolID, StateID>> StartStates;
+    return Table;
+  }
 };
 
-LRTable LRTable::buildForTests(const GrammarTable &GT,
-                               llvm::ArrayRef<Entry> Entries) {
+LRTable LRTable::buildForTests(const Grammar &G, llvm::ArrayRef<Entry> Entries,
+                               llvm::ArrayRef<ReduceEntry> Reduces) {
   StateID MaxState = 0;
   for (const auto &Entry : Entries) {
     MaxState = std::max(MaxState, Entry.State);
     if (Entry.Act.kind() == LRTable::Action::Shift)
       MaxState = std::max(MaxState, Entry.Act.getShiftState());
     if (Entry.Act.kind() == LRTable::Action::GoTo)
       MaxState = std::max(MaxState, Entry.Act.getGoToState());
   }
-  Builder Build({});
-  for (const Entry &E : Entries)
-    Build.insert(E);
-  return std::move(Build).build(GT, /*NumStates=*/MaxState + 1);
+  Builder Build;
+  Build.Entries.insert(Entries.begin(), Entries.end());
+  for (const ReduceEntry &E : Reduces)
+    Build.Reduces[E.State].insert(E.Rule);
+  Build.FollowSets = followSets(G);
+  return std::move(Build).build(/*NumStates=*/MaxState + 1);
 }
 
 LRTable LRTable::buildSLR(const Grammar &G) {
   auto Graph = LRGraph::buildLR0(G);
-  Builder Build(Graph.startStates());
+  Builder Build;
+  Build.StartStates = Graph.startStates();
   for (const auto &T : Graph.edges()) {
     Action Act = isToken(T.Label) ? Action::shift(T.Dst) : Action::goTo(T.Dst);
-    Build.insert({T.Src, T.Label, Act});
+    Build.Entries.insert({T.Src, T.Label, Act});
   }
+  Build.FollowSets = followSets(G);
   assert(Graph.states().size() <= (1 << StateBits) &&
          "Graph states execceds the maximum limit!");
-  auto FollowSets = followSets(G);
+  // Add reduce actions.
   for (StateID SID = 0; SID < Graph.states().size(); ++SID) {
     for (const Item &I : Graph.states()[SID].Items) {
       // If we've just parsed the start symbol, this means we successfully parse
       // the input. We don't add the reduce action of `_ := start_symbol` in the
       // LRTable (the GLR parser handles it specifically).
       if (G.lookupRule(I.rule()).Target == G.underscore() && !I.hasNext())
         continue;
-      if (!I.hasNext()) {
+      if (!I.hasNext())
         // If we've reached the end of a rule A := ..., then we can reduce if
         // the next token is in the follow set of A.
-        for (SymbolID Follow : FollowSets[G.lookupRule(I.rule()).Target]) {
-          assert(isToken(Follow));
-          Build.insert({SID, Follow, Action::reduce(I.rule())});
-        }
-      }
+        Build.Reduces[SID].insert(I.rule());
     }
   }
-  return std::move(Build).build(G.table(), Graph.states().size());
+  return std::move(Build).build(Graph.states().size());
 }
 
 } // namespace pseudo
 } // namespace clang

clang-tools-extra/pseudo/test/lr-build-basic.test

 # GRAPH-NEXT: State 2
 # GRAPH-NEXT:     expr := id • 
 # GRAPH-NEXT: State 3
 # GRAPH-NEXT:     id := IDENTIFIER • 
 
 # RUN: clang-pseudo -grammar %s -print-table | FileCheck %s --check-prefix=TABLE
 #      TABLE: LRTable:
 # TABLE-NEXT: State 0
-# TABLE-NEXT:     'IDENTIFIER': shift state 3
-# TABLE-NEXT:     'expr': go to state 1
-# TABLE-NEXT:     'id': go to state 2
+# TABLE-NEXT:     IDENTIFIER: shift state 3
+# TABLE-NEXT:     expr: go to state 1
+# TABLE-NEXT:     id: go to state 2
 # TABLE-NEXT: State 1
 # TABLE-NEXT: State 2
-# TABLE-NEXT:     'EOF': reduce by rule 1 'expr := id'
+# TABLE-NEXT:     EOF: reduce by rule 1 'expr := id'
 # TABLE-NEXT: State 3
-# TABLE-NEXT:     'EOF': reduce by rule 0 'id := IDENTIFIER'
+# TABLE-NEXT:     EOF: reduce by rule 0 'id := IDENTIFIER'

clang-tools-extra/pseudo/test/lr-build-conflicts.test

 # GRAPH-NEXT: 1 ->[-] 3
 # GRAPH-NEXT: 3 ->[expr] 4
 # GRAPH-NEXT: 3 ->[IDENTIFIER] 2
 # GRAPH-NEXT: 4 ->[-] 3
 
 # RUN: clang-pseudo -grammar %s -print-table | FileCheck %s --check-prefix=TABLE
 #      TABLE: LRTable:
 # TABLE-NEXT: State 0
-# TABLE-NEXT:     'IDENTIFIER': shift state 2
-# TABLE-NEXT:     'expr': go to state 1
+# TABLE-NEXT:     IDENTIFIER: shift state 2
+# TABLE-NEXT:     expr: go to state 1
 # TABLE-NEXT: State 1
-# TABLE-NEXT:     '-': shift state 3
+# TABLE-NEXT:     -: shift state 3
 # TABLE-NEXT: State 2
-# TABLE-NEXT:     'EOF': reduce by rule 1 'expr := IDENTIFIER'
-# TABLE-NEXT:     '-': reduce by rule 1 'expr := IDENTIFIER'
+# TABLE-NEXT:     EOF -: reduce by rule 1 'expr := IDENTIFIER'
 # TABLE-NEXT: State 3
-# TABLE-NEXT:     'IDENTIFIER': shift state 2
-# TABLE-NEXT:     'expr': go to state 4
+# TABLE-NEXT:     IDENTIFIER: shift state 2
+# TABLE-NEXT:     expr: go to state 4
 # TABLE-NEXT: State 4
-# TABLE-NEXT:     'EOF': reduce by rule 0 'expr := expr - expr'
-# TABLE-NEXT:     '-': shift state 3
-# TABLE-NEXT:     '-': reduce by rule 0 'expr := expr - expr'
+# TABLE-NEXT:     -: shift state 3
+# TABLE-NEXT:     EOF -: reduce by rule 0 'expr := expr - expr'

clang-tools-extra/pseudo/unittests/GLRTest.cpp

for (const auto *Head : Heads)
     OS << *Head << "\n";
   return OS;
 }
 
 namespace {
 
 using Action = LRTable::Action;
 using testing::AllOf;
+using testing::ElementsAre;
 using testing::UnorderedElementsAre;
 
 MATCHER_P(state, StateID, "") { return arg->State == StateID; }
 MATCHER_P(parsedSymbol, FNode, "") { return arg->Payload == FNode; }
 MATCHER_P(parsedSymbolID, SID, "") { return arg->Payload->symbol() == SID; }
 
 testing::Matcher<const GSS::Node *>
 parents(llvm::ArrayRef<const GSS::Node *> Parents) {
auto *GSSNode1 = GSStack.addNode(/*State=*/1, /*ForestNode=*/nullptr,
                                    /*Parents=*/{GSSNode0});
   auto *GSSNode2 = GSStack.addNode(/*State=*/2, /*ForestNode=*/nullptr,
                                    /*Parents=*/{GSSNode0});
   auto *GSSNode3 = GSStack.addNode(/*State=*/3, /*ForestNode=*/nullptr,
                                    /*Parents=*/{GSSNode0});
 
   buildGrammar({}, {}); // Create a fake empty grammar.
   LRTable T =
-      LRTable::buildForTests(G->table(), /*Entries=*/{
+      LRTable::buildForTests(*G, /*Entries=*/
+                             {
                                  {1, tokenSymbol(tok::semi), Action::shift(4)},
                                  {2, tokenSymbol(tok::semi), Action::shift(4)},
                                  {3, tokenSymbol(tok::semi), Action::shift(5)},
-                             });
+                             },
+                             {});
 
   ForestNode &SemiTerminal = Arena.createTerminal(tok::semi, 0);
   std::vector<const GSS::Node *> NewHeads;
   glrShift({GSSNode1, GSSNode2, GSSNode3}, SemiTerminal,
            {*G, T, Arena, GSStack}, NewHeads);
 
   EXPECT_THAT(NewHeads,
               UnorderedElementsAre(AllOf(state(4), parsedSymbol(&SemiTerminal),
TEST_F(GLRTest, ReduceConflictsSplitting) {
   //  After reducing 1 by `class-name := IDENTIFIER` and
   //                      `enum-name := IDENTIFIER`:
   //    0--2(class-name)    // 2 is goto(0, class-name)
   //    └--3(enum-name)     // 3 is goto(0, enum-name)
   buildGrammar({"class-name", "enum-name"},
                {"class-name := IDENTIFIER", "enum-name := IDENTIFIER"});
 
   LRTable Table = LRTable::buildForTests(
-      G->table(), {
-                      {/*State=*/0, id("class-name"), Action::goTo(2)},
-                      {/*State=*/0, id("enum-name"), Action::goTo(3)},
-                      {/*State=*/1, tokenSymbol(tok::l_brace),
-                       Action::reduce(ruleFor("class-name"))},
-                      {/*State=*/1, tokenSymbol(tok::l_brace),
-                       Action::reduce(ruleFor("enum-name"))},
-                  });
+      *G,
+      {
+      },
+      {
+          {/*State=*/1, ruleFor("class-name")},
+          {/*State=*/1, ruleFor("enum-name")},
 
   const auto *GSSNode0 =
       GSStack.addNode(/*State=*/0, /*ForestNode=*/nullptr, /*Parents=*/{});
   const auto *GSSNode1 =
       GSStack.addNode(1, &Arena.createTerminal(tok::identifier, 0), {GSSNode0});
 
   std::vector<const GSS::Node *> Heads = {GSSNode1};
-  glrReduce(Heads, tokenSymbol(tok::l_brace), {*G, Table, Arena, GSStack});
+  glrReduce(Heads, tokenSymbol(tok::eof), {*G, Table, Arena, GSStack});
   EXPECT_THAT(Heads, UnorderedElementsAre(
                          GSSNode1,
                          AllOf(state(2), parsedSymbolID(id("class-name")),
                                parents({GSSNode0})),
                          AllOf(state(3), parsedSymbolID(id("enum-name")),
                                parents({GSSNode0}))))
       << Heads;
 }
const auto *GSSNode2 =
       GSStack.addNode(/*State=*/2, /*ForestNode=*/ClassNameNode, /*Parents=*/{});
   const auto *GSSNode3 =
       GSStack.addNode(/*State=*/3, /*ForestNode=*/EnumNameNode, /*Parents=*/{});
   const auto *GSSNode4 = GSStack.addNode(
       /*State=*/4, &Arena.createTerminal(tok::star, /*TokenIndex=*/1),
       /*Parents=*/{GSSNode2, GSSNode3});
 
   LRTable Table = LRTable::buildForTests(
-      G->table(),
+      *G,
       {
           {/*State=*/2, id("ptr-operator"), Action::goTo(/*NextState=*/5)},
           {/*State=*/3, id("ptr-operator"), Action::goTo(/*NextState=*/6)},
-          {/*State=*/4, tokenSymbol(tok::identifier),
-           Action::reduce(ruleFor("ptr-operator"))},
+      },
+      {
+          {/*State=*/4, ruleFor("ptr-operator")},
       });
   std::vector<const GSS::Node *> Heads = {GSSNode4};
-  glrReduce(Heads, tokenSymbol(tok::identifier), {*G, Table, Arena, GSStack});
+  glrReduce(Heads, tokenSymbol(tok::eof), {*G, Table, Arena, GSStack});
 
   EXPECT_THAT(Heads, UnorderedElementsAre(
                          GSSNode4,
                          AllOf(state(5), parsedSymbolID(id("ptr-operator")),
                                parents({GSSNode2})),
                          AllOf(state(6), parsedSymbolID(id("ptr-operator")),
                                parents({GSSNode3}))))
       << Heads;
const auto *GSSNode3 =
       GSStack.addNode(/*State=*/3, /*ForestNode=*/ClassNameNode,
                       /*Parents=*/{GSSNode1});
   const auto *GSSNode4 =
       GSStack.addNode(/*State=*/4, /*ForestNode=*/EnumNameNode,
                       /*Parents=*/{GSSNode2});
 
   // FIXME: figure out a way to get rid of the hard-coded reduce RuleID!
   LRTable Table = LRTable::buildForTests(
-      G->table(),
+      *G,
       {
           {/*State=*/1, id("type-name"), Action::goTo(/*NextState=*/5)},
           {/*State=*/2, id("type-name"), Action::goTo(/*NextState=*/5)},
-          {/*State=*/3, tokenSymbol(tok::l_paren),
-           Action::reduce(/* type-name := class-name */ 0)},
-          {/*State=*/4, tokenSymbol(tok::l_paren),
-           Action::reduce(/* type-name := enum-name */ 1)},
+      },
+      {
+          {/*State=*/3, /* type-name := class-name */ 0},
+          {/*State=*/4, /* type-name := enum-name */ 1},
       });
   std::vector<const GSS::Node *> Heads = {GSSNode3, GSSNode4};
-  glrReduce(Heads, tokenSymbol(tok::l_paren), {*G, Table, Arena, GSStack});
+  glrReduce(Heads, tokenSymbol(tok::eof), {*G, Table, Arena, GSStack});
 
   // Verify that the stack heads are joint at state 5 after reduces.
   EXPECT_THAT(Heads, UnorderedElementsAre(GSSNode3, GSSNode4,
                                           AllOf(state(5),
                                                 parsedSymbolID(id("type-name")),
                                                 parents({GSSNode1, GSSNode2}))))
       << Heads;
   // Verify that we create an ambiguous ForestNode of two parses of `type-name`.
TEST_F(GLRTest, ReduceJoiningWithSameBase) {
   const auto *GSSNode3 =
       GSStack.addNode(/*State=*/3, /*ForestNode=*/StartTerminal,
                       /*Parents=*/{GSSNode1});
   const auto *GSSNode4 =
       GSStack.addNode(/*State=*/4, /*ForestNode=*/StartTerminal,
                       /*Parents=*/{GSSNode2});
 
   // FIXME: figure out a way to get rid of the hard-coded reduce RuleID!
-  LRTable Table = LRTable::buildForTests(
-      G->table(), {
+  LRTable Table =
+      LRTable::buildForTests(*G,
+                             {
                                  {/*State=*/0, id("pointer"), Action::goTo(5)},
-                      {3, tokenSymbol(tok::l_paren),
-                       Action::reduce(/* pointer := class-name */ 0)},
-                      {4, tokenSymbol(tok::l_paren),
-                       Action::reduce(/* pointer := enum-name */ 1)},
+                             },
+                             {
+                                 {3, /* pointer := class-name */ 0},
+                                 {4, /* pointer := enum-name */ 1},
                              });
   std::vector<const GSS::Node *> Heads = {GSSNode3, GSSNode4};
-  glrReduce(Heads, tokenSymbol(tok::l_paren), {*G, Table, Arena, GSStack});
+  glrReduce(Heads, tokenSymbol(tok::eof), {*G, Table, Arena, GSStack});
 
   EXPECT_THAT(
       Heads, UnorderedElementsAre(GSSNode3, GSSNode4,
                                   AllOf(state(5), parsedSymbolID(id("pointer")),
                                         parents({GSSNode0}))))
       << Heads;
   EXPECT_EQ(Heads.back()->Payload->dumpRecursive(*G),
             "[  0, end) pointer := <ambiguous>\n"
             "[  0, end) ├─pointer := class-name *\n"
             "[  0,   1) │ ├─class-name := <opaque>\n"
             "[  1, end) │ └─* := tok[1]\n"
             "[  0, end) └─pointer := enum-name *\n"
             "[  0,   1)   ├─enum-name := <opaque>\n"
             "[  1, end)   └─* := tok[1]\n");
 }
 
+TEST_F(GLRTest, ReduceLookahead) {
+  // A term can be followed by +, but not by -.
+  buildGrammar({"sum", "term"}, {"expr := term + term", "term := IDENTIFIER"});
+  LRTable Table =
+      LRTable::buildForTests(*G,
+                             {
+                                 {/*State=*/0, id("term"), Action::goTo(2)},
+                             },
+                             {
+                                 {/*State=*/1, 0},
+                             });
+
+  auto *Identifier = &Arena.createTerminal(tok::identifier, /*Start=*/0);
+
+  const auto *Root =
+      GSStack.addNode(/*State=*/0, /*ForestNode=*/nullptr, /*Parents=*/{});
+  const auto *GSSNode1 =
+      GSStack.addNode(/*State=*/1, /*ForestNode=*/Identifier, {Root});
+
+  // When the lookahead is +, reduce is performed.
+  std::vector<const GSS::Node *> Heads = {GSSNode1};
+  glrReduce(Heads, tokenSymbol(tok::plus), {*G, Table, Arena, GSStack});
+  EXPECT_THAT(Heads,
+              ElementsAre(GSSNode1, AllOf(state(2), parsedSymbolID(id("term")),
+                                          parents(Root))));
+
+  // When the lookahead is -, reduce is not performed.
+  Heads = {GSSNode1};
+  glrReduce(Heads, tokenSymbol(tok::minus), {*G, Table, Arena, GSStack});
+  EXPECT_THAT(Heads, ElementsAre(GSSNode1));
+}
+
 TEST_F(GLRTest, PerfectForestNodeSharing) {
   // Run the GLR on a simple grammar and test that we build exactly one forest
   // node per (SymbolID, token range).
 
   // This is a grmammar where the original parsing-stack-based forest node
   // sharing approach will fail. In its LR0 graph, it has two states containing
   // item `expr := • IDENTIFIER`, and both have different goto states on the
   // nonterminal `expr`.

clang-tools-extra/pseudo/unittests/LRTableTest.cpp

 //===--- LRTableTest.cpp - ---------------------------------------*- C++-*-===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 
 #include "clang-pseudo/grammar/LRTable.h"
 #include "clang-pseudo/grammar/Grammar.h"
 #include "clang/Basic/TokenKinds.h"
+#include "llvm/Testing/Support/SupportHelpers.h"
 #include "gmock/gmock.h"
 #include "gtest/gtest.h"
 #include <vector>
 
 namespace clang {
 namespace pseudo {
 namespace {
 
-using testing::IsEmpty;
-using testing::UnorderedElementsAre;
+using llvm::ValueIs;
+using testing::ElementsAre;
 using Action = LRTable::Action;
 
 TEST(LRTable, Builder) {
-  GrammarTable GTable;
-
-  //           eof   semi  ...
-  // +-------+----+-------+---
-  // |state0 |    | s0,r0 |...
-  // |state1 | acc|       |...
-  // |state2 |    |  r1   |...
-  // +-------+----+-------+---
+  std::vector<std::string> GrammarDiags;
+  auto G = Grammar::parseBNF(R"bnf(
+    _ := expr            # rule 0
+    expr := term         # rule 1
+    expr := expr + term  # rule 2
+    term := IDENTIFIER   # rule 3
+  )bnf",
+                             GrammarDiags);
+  EXPECT_THAT(GrammarDiags, testing::IsEmpty());
+
+  SymbolID Term = *G->findNonterminal("term");
+  SymbolID Eof = tokenSymbol(tok::eof);
+  SymbolID Identifier = tokenSymbol(tok::identifier);
+  SymbolID Plus = tokenSymbol(tok::plus);
+
+  //           eof  IDENT   term
+  // +-------+----+-------+------+
+  // |state0 |    | s0    |      |
+  // |state1 |    |       | g3   |
+  // |state2 |    |       |      |
+  // +-------+----+-------+------+-------
   std::vector<LRTable::Entry> Entries = {
-      {/* State */ 0, tokenSymbol(tok::semi), Action::shift(0)},
-      {/* State */ 0, tokenSymbol(tok::semi), Action::reduce(0)},
-      {/* State */ 1, tokenSymbol(tok::eof), Action::reduce(2)},
-      {/* State */ 2, tokenSymbol(tok::semi), Action::reduce(1)}};
-  GrammarTable GT;
-  LRTable T = LRTable::buildForTests(GT, Entries);
-  EXPECT_THAT(T.find(0, tokenSymbol(tok::eof)), IsEmpty());
-  EXPECT_THAT(T.find(0, tokenSymbol(tok::semi)),
-              UnorderedElementsAre(Action::shift(0), Action::reduce(0)));
-  EXPECT_THAT(T.find(1, tokenSymbol(tok::eof)),
-              UnorderedElementsAre(Action::reduce(2)));
-  EXPECT_THAT(T.find(1, tokenSymbol(tok::semi)), IsEmpty());
-  EXPECT_THAT(T.find(2, tokenSymbol(tok::semi)),
-              UnorderedElementsAre(Action::reduce(1)));
+      {/* State */ 0, Identifier, Action::shift(0)},
+      {/* State */ 1, Term, Action::goTo(3)},
+  };
+  std::vector<LRTable::ReduceEntry> ReduceEntries = {
+      {/*State=*/0, /*Rule=*/0},
+      {/*State=*/1, /*Rule=*/2},
+      {/*State=*/2, /*Rule=*/1},
+  };
+  LRTable T = LRTable::buildForTests(*G, Entries, ReduceEntries);
+  EXPECT_EQ(T.getShiftState(0, Eof), llvm::None);
+  EXPECT_THAT(T.getShiftState(0, Identifier), ValueIs(0));
+  EXPECT_THAT(T.getReduceRules(0), ElementsAre(0));
+
+  EXPECT_EQ(T.getShiftState(1, Eof), llvm::None);
+  EXPECT_EQ(T.getShiftState(1, Identifier), llvm::None);
+  EXPECT_EQ(T.getGoToState(1, Term), 3);
+  EXPECT_THAT(T.getReduceRules(1), ElementsAre(2));
+
   // Verify the behaivor for other non-available-actions terminals.
-  EXPECT_THAT(T.find(2, tokenSymbol(tok::kw_int)), IsEmpty());
+  SymbolID Int = tokenSymbol(tok::kw_int);
+  EXPECT_EQ(T.getShiftState(2, Int), llvm::None);
+
+  // Check follow sets.
+  EXPECT_TRUE(T.canFollow(Term, Plus));
+  EXPECT_TRUE(T.canFollow(Term, Eof));
+  EXPECT_FALSE(T.canFollow(Term, Int));
 }
 
 } // namespace
 } // namespace pseudo
 } // namespace clang