[PoC] Support Token Categories

lib/lrama.rb

@@ -1,6 +1,7 @@
 # frozen_string_literal: true
 
 require_relative "lrama/bitmap"
+require_relative "lrama/category"
 require_relative "lrama/command"
 require_relative "lrama/context"
 require_relative "lrama/counterexamples"

lib/lrama/category.rb

@@ -0,0 +1,27 @@
+# rbs_inline: enabled
+# frozen_string_literal: true
+
+module Lrama
+  class Category
+    attr_reader :id #: Lrama::Lexer::Token::Ident
+    attr_reader :name #: String
+    attr_reader :tokens #: Array[Lrama::Lexer::Token::Ident]
+    attr_reader :tag #: Lexer::Token::Tag
+
+    # @rbs (id: Lrama::Lexer::Token::Ident) -> void
+    def initialize(id:)
+      @id = id
+      @name = id.s_value
+      @tokens = []
+      @tag = nil
+    end
+
+    # @rbs (Array[Lrama::Lexer::Token::Ident] tokens, Lexer::Token::Tag tag) -> void
+    def add_tokens(tokens, tag)
+      @tag = tag
+      tokens.each do |token|
+        @tokens << Lrama::Lexer::Token::Ident.new(s_value: token.s_value, location: token.location)
+      end
+    end
+  end
+end

lib/lrama/grammar.rb

@@ -60,6 +60,7 @@ def initialize(rule_counter, define = {})
       @locations = false
       @define = define
       @required = false
+      @categories = {}
 
       append_special_symbols
     end
@@ -113,6 +114,14 @@ def set_union(code, lineno)
       @union = Union.new(code: code, lineno: lineno)
     end
 
+    def add_category(id:, tokens:, tag:)
+      unless category = find_category(id.s_value)
+        category = Lrama::Category.new(id: id)
+        @categories[id.s_value] = category
+      end
+      category.add_tokens(tokens, tag)
+    end
+
     def add_rule_builder(builder)
       @rule_builders << builder
     end
@@ -182,6 +191,10 @@ def unique_rule_s_values
       @rules.map {|rule| rule.lhs.id.s_value }.uniq
     end
 
+    def find_category(name)
+      @categories[name]
+    end
+
     def ielr_defined?
       @define.key?('lr.type') && @define['lr.type'] == 'ielr'
     end
@@ -335,9 +348,19 @@ def normalize_rules
         end
       end
 
+      expand_categories
+
       @rules.sort_by!(&:id)
     end
 
+    def expand_categories
+      @categories.values.each do |category|
+        lhs = Lrama::Lexer::Token::Ident.new(s_value: category.name, location: category.id.location)
+        add_nterm(id: lhs, tag: category.tag)
+        @rules << Rule.new(id: @rule_counter.increment, _lhs: lhs, _rhs: category.tokens, lhs_tag: category.tag, token_code: nil, lineno: lhs.line)
+      end
+    end
+
     # Collect symbols from rules
     def collect_symbols
       @rules.flat_map(&:_rhs).each do |s|

lib/lrama/lexer.rb

@@ -14,6 +14,7 @@ class Lexer
     SYMBOLS = ['%{', '%}', '%%', '{', '}', '\[', '\]', '\(', '\)', '\,', ':', '\|', ';'].freeze
     PERCENT_TOKENS = %w(
       %union
+      %token-categories
       %token
       %type
       %nterm
@@ -42,7 +43,6 @@ class Lexer
       %no-stdlib
       %inline
       %locations
-      %categories
     ).freeze
 
     def initialize(grammar_file)

parser.y

@@ -167,6 +167,7 @@ rule
                           }
                           @precedence_number += 1
                         }
+                    | "%token-categories" token_categories_declarations
 
   token_declarations: TAG? token_declaration+
                         {
@@ -183,6 +184,11 @@ rule
 
   token_declaration: id INTEGER? alias { result = val }
 
+  token_categories_declarations: TAG? IDENT_COLON ":" id+
+                                {
+                                  @grammar.add_category(id: val[1], tokens: val[3], tag: val[0])
+                                }
+
   rule_declaration: "%rule" IDENTIFIER "(" rule_args ")" TAG? ":" rule_rhs_list
                       {
                         rule = Grammar::Parameterized::Rule.new(val[1].s_value, val[3], val[7], tag: val[5])

sig/generated/lrama/category.rbs

@@ -0,0 +1,19 @@
+# Generated from lib/lrama/category.rb with RBS::Inline
+
+module Lrama
+  class Category
+    attr_reader id: Lrama::Lexer::Token::Ident
+
+    attr_reader name: String
+
+    attr_reader tokens: Array[Lrama::Lexer::Token::Ident]
+
+    attr_reader tag: Lexer::Token::Tag
+
+    # @rbs (id: Lrama::Lexer::Token::Ident) -> void
+    def initialize: (id: Lrama::Lexer::Token::Ident) -> void
+
+    # @rbs (Array[Lrama::Lexer::Token::Ident] tokens, Lexer::Token::Tag tag) -> void
+    def add_tokens: (Array[Lrama::Lexer::Token::Ident] tokens, Lexer::Token::Tag tag) -> void
+  end
+end

sig/lrama/grammar.rbs

@@ -21,6 +21,9 @@ module Lrama
     @aux: Auxiliary
     @no_stdlib: bool
     @locations: bool
+    @define: Hash[String, String]
+    @required: bool
+    @categories: Hash[String, Lrama::Category]
     @union: Union
 
     extend Forwardable
@@ -63,6 +66,7 @@ module Lrama
     def add_precedence: (Grammar::Symbol sym, Integer precedence) -> Precedence
     def set_precedence: (Grammar::Symbol sym, Precedence precedence) -> (Precedence | bot)
     def set_union: (Grammar::Code::NoReferenceCode code, Integer lineno) -> Union
+    def add_category: (id: Lrama::Lexer::Token::Ident, tokens: Array[Lrama::Lexer::Token::Ident], tag: Lexer::Token::Tag) -> void
     def add_rule_builder: (RuleBuilder builder) -> Array[RuleBuilder]
     def add_parameterized_rule: (Parameterized::Rule rule) -> Array[Parameterized::Rule]
     def parameterized_rules: () -> Array[Parameterized::Rule]
@@ -75,6 +79,7 @@ module Lrama
     def validate!: () -> void
     def find_rules_by_symbol!: (Grammar::Symbol sym) -> Array[Rule]
     def find_rules_by_symbol: (Grammar::Symbol sym) -> Array[Rule]?
+    def find_category: (String) -> Lrama::Category
 
     private
 
@@ -84,6 +89,7 @@ module Lrama
     def setup_rules: () -> Array[RuleBuilder]
     def append_special_symbols: () -> Grammar::Symbol
     def normalize_rules: () -> Array[Rule]
+    def expand_categories: () -> void
     def collect_symbols: () -> Array[Lexer::Token]
     def set_lhs_and_rhs: () -> void
     def fill_default_precedence: () -> void

spec/fixtures/categories/basic.y

@@ -0,0 +1,41 @@
+/*
+ * This is comment for this file.
+ */
+
+%{
+// Prologue
+static int yylex(YYSTYPE *val, YYLTYPE *loc);
+static int yyerror(YYLTYPE *loc, const char *str);
+%}
+
+%union {
+    int i;
+    char* s;
+}
+
+%token <i> integer number count
+%token <s> string identifier
+%token-categories num: integer number count
+%token-categories str: string identifier
+
+%%
+
+program         : num
+                | str
+                ;
+
+%%
+
+static int yylex(YYSTYPE *yylval, YYLTYPE *loc)
+{
+  return 0;
+}
+
+static int yyerror(YYLTYPE *loc, const char *str)
+{
+  return 0;
+}
+
+int main(int argc, char *argv[])
+{
+}

spec/lrama/parser_spec.rb

@@ -2869,6 +2869,87 @@
       end
     end
 
+    context "when categories" do
+      let(:grammar) do
+        grammar = Lrama::Parser.new(y, path).parse
+        grammar.prepare
+        grammar.validate!
+        grammar
+      end
+      let(:y) { File.read(fixture_path(path)) }
+      let(:path) { "categories/basic.y" }
+
+      it "expands categories" do
+        expect(grammar.nterms.sort_by(&:number)).to match_symbols([
+          Sym.new(id: T::Ident.new(s_value: "$accept"), alias_name: nil, number: 8, tag: nil, term: false, token_id: 0, nullable: false),
+          Sym.new(id: T::Ident.new(s_value: "program"), alias_name: nil, number: 9, tag: nil, term: false, token_id: 1, nullable: false),
+          Sym.new(id: T::Ident.new(s_value: "num"), alias_name: nil, number: 10, tag: nil, term: false, token_id: 2, nullable: false),
+          Sym.new(id: T::Ident.new(s_value: "str"), alias_name: nil, number: 11, tag: nil, term: false, token_id: 3, nullable: false),
+        ])
+
+        expect(grammar.rules).to eq([
+          Rule.new(
+            id: 0,
+            lhs: grammar.find_symbol_by_s_value!("$accept"),
+            rhs: [
+              grammar.find_symbol_by_s_value!("program"),
+              grammar.find_symbol_by_s_value!("YYEOF"),
+            ],
+            token_code: nil,
+            nullable: false,
+            precedence_sym: grammar.find_symbol_by_s_value!("YYEOF"),
+            lineno: 23,
+          ),
+          Rule.new(
+            id: 1,
+            lhs: grammar.find_symbol_by_s_value!("program"),
+            rhs: [
+              grammar.find_symbol_by_s_value!("num"),
+            ],
+            token_code: nil,
+            nullable: false,
+            precedence_sym: nil,
+            lineno: 23,
+          ),
+          Rule.new(
+            id: 2,
+            lhs: grammar.find_symbol_by_s_value!("program"),
+            rhs: [
+              grammar.find_symbol_by_s_value!("str"),
+            ],
+            token_code: nil,
+            nullable: false,
+            precedence_sym: nil,
+            lineno: 24,
+          ),
+          Rule.new(
+            id: 3,
+            lhs: grammar.find_symbol_by_s_value!("num"),
+            rhs: [
+              grammar.find_symbol_by_s_value!("integer"),
+              grammar.find_symbol_by_s_value!("number"),
+              grammar.find_symbol_by_s_value!("count"),
+            ],
+            token_code: nil,
+            nullable: false,
+            precedence_sym: grammar.find_symbol_by_s_value!("count"),
+            lineno: 18,
+          ),
+          Rule.new(
+            id: 4,
+            lhs: grammar.find_symbol_by_s_value!("str"),
+            rhs: [
+              grammar.find_symbol_by_s_value!("string"),
+              grammar.find_symbol_by_s_value!("identifier"),
+            ],
+            nullable: false,
+            precedence_sym: grammar.find_symbol_by_s_value!("identifier"),
+            lineno: 19,
+          )
+        ])
+      end
+    end
+
     it "; for rules is optional" do
       y = header + <<~INPUT
         %%