Token locations

Rakefile

@@ -2,6 +2,8 @@ require "bundler/gem_tasks"
 require "rake/testtask"
 require "rbconfig"
 
+$LOAD_PATH << File.join(__dir__, "test")
+
 ruby = ENV["RUBY"] || RbConfig.ruby
 racc = ENV.fetch("RACC", "racc")
 rbs = File.join(__dir__, "exe/rbs")

lib/rbs/ast/declarations.rb

@@ -7,7 +7,7 @@ class Base
       class ModuleTypeParams
         attr_reader :params
 
-        TypeParam = _ = Struct.new(:name, :variance, :skip_validation, keyword_init: true) do
+        TypeParam = _ = Struct.new(:name, :variance, :skip_validation, :location, keyword_init: true) do
           # @implements TypeParam
 
           def to_json(state = _ = nil)
@@ -17,6 +17,19 @@ def to_json(state = _ = nil)
               skip_validation: skip_validation,
             }.to_json(state)
           end
+
+          def ==(other)
+            other.is_a?(TypeParam) &&
+              other.name == name &&
+              other.variance == variance &&
+              other.skip_validation == skip_validation
+          end
+
+          alias eql? ==
+
+          def hash
+            self.class.hash ^ name.hash ^ variance.hash ^ skip_validation.hash
+          end
         end
 
         def initialize()
@@ -82,7 +95,7 @@ def rename_to(names)
           ModuleTypeParams.new().tap do |params|
             names.each.with_index do |new_name, index|
               param = self.params[index]
-              params.add(TypeParam.new(name: new_name, variance: param.variance, skip_validation: param.skip_validation))
+              params.add(TypeParam.new(name: new_name, variance: param.variance, skip_validation: param.skip_validation, location: param.location))
             end
           end
         end

lib/rbs/location.rb

@@ -42,6 +42,10 @@ def end_loc
       @end_loc ||= buffer.pos_to_loc(end_pos)
     end
 
+    def range
+      start_pos...end_pos
+    end
+
     def source
       @source ||= buffer.content[start_pos...end_pos] or raise
     end
@@ -109,5 +113,105 @@ def to_json(state = _ = nil)
         }
       }.to_json(state)
     end
+
+    def with_children(required: {}, optional: {})
+      # @type var required: Hash[Symbol, Range[Integer] | Location]
+      # @type var optional: Hash[Symbol, Range[Integer] | Location | nil]
+
+      this = WithChildren.new(buffer: buffer, start_pos: start_pos, end_pos: end_pos)
+
+      req = required.transform_values do |value|
+        case value
+        when Location
+          value.range
+        else
+          value
+        end
+      end
+
+      opt = optional.transform_values do |value|
+        case value
+        when Location
+          value.range
+        else
+          value
+        end
+      end
+
+      this.required_children.merge!(req)
+      this.optional_children.merge!(opt)
+
+      this
+    end
+
+    class WithChildren < Location
+      attr_reader :required_children, :optional_children
+
+      def initialize(buffer:, start_pos:, end_pos:)
+        super(buffer: buffer, start_pos: start_pos, end_pos: end_pos)
+
+        @optional_children = {}
+        @required_children = {}
+      end
+
+      def initialize_copy(from)
+        required_children.merge!(from.required_children)
+        optional_children.merge!(from.optional_children)
+        self
+      end
+
+      def [](key)
+        case
+        when required_children.key?(_ = key)
+          range = required_children[_ = key]
+          Location.new(buffer: buffer, start_pos: range.begin, end_pos: range.end)
+        when optional_children.key?(_ = key)
+          range = required_children[_ = key] || optional_children[_ = key]
+          if range
+            Location.new(buffer: buffer, start_pos: range.begin, end_pos: range.end)
+          end
+        else
+          raise "Unknown key given: `#{key}`"
+        end
+      end
+
+      def merge_required(hash)
+        this = dup
+
+        h = hash.transform_values do |value|
+          case value
+          when Range
+            value
+          when Location
+            value.range
+          else
+            raise
+          end
+        end
+
+        this.required_children.merge!(h)
+
+        this
+      end
+
+      def merge_optional(hash)
+        this = dup
+
+        h = hash.transform_values do |value|
+          case value
+          when Range
+            value
+          when Location
+            value.range
+          else
+            nil
+          end
+        end
+
+        this.optional_children.merge!(h)
+
+        this
+      end
+    end
   end
 end