Let Locator support use directives

lib/rbs/locator.rb

@@ -2,21 +2,24 @@
 
 module RBS
   class Locator
-    attr_reader :decls
+    attr_reader :decls, :dirs, :buffer
 
-    def initialize(decls:)
+    def initialize(buffer:, dirs:, decls:)
+      @buffer = buffer
+      @dirs = dirs
       @decls = decls
     end
 
-    def buffer
-      decls[0].location&.buffer or raise
-    end
-
     def find(line:, column:)
       pos = buffer.loc_to_pos([line, column])
 
+      dirs.each do |dir|
+        array = [] #: Array[component]
+        find_in_directive(pos, dir, array) and return array
+      end
+
       decls.each do |decl|
-        array = []
+        array = [] #: Array[component]
         find_in_decl(pos, decl: decl, array: array) and return array
       end
 
@@ -36,6 +39,22 @@ def find2(line:, column:)
       end
     end
 
+    def find_in_directive(pos, dir, array)
+      if test_loc(pos, location: dir.location)
+        array.unshift(dir)
+
+        dir.clauses.each do |clause|
+          if test_loc(pos, location: clause.location)
+            array.unshift(clause)
+            find_in_loc(pos, location: clause.location, array: array)
+            return true
+          end
+        end
+      end
+
+      false
+    end
+
     def find_in_decl(pos, decl:, array:)
       if test_loc(pos, location: decl.location)
         array.unshift(decl)

sig/locator.rbs

@@ -10,14 +10,22 @@ module RBS
                    | AST::TypeParam
                    | AST::Declarations::Class::Super
                    | AST::Declarations::Module::Self
+                   | AST::Directives::t
+                   | AST::Directives::Use::clause
+
+    # The buffer the location points to
+    #
+    attr_reader buffer: Buffer
 
     # Array of _top-level_ declarations.
     #
     attr_reader decls: Array[AST::Declarations::t]
 
-    def initialize: (decls: Array[AST::Declarations::t]) -> void
+    # Array of directives.
+    #
+    attr_reader dirs: Array[AST::Directives::t]
 
-    def buffer: () -> Buffer
+    def initialize: (buffer: Buffer, decls: Array[AST::Declarations::t], dirs: Array[AST::Directives::t]) -> void
 
     # Returns list of components.
     # Inner component comes first.
@@ -29,6 +37,10 @@ module RBS
     #
     def find2: (line: Integer, column: Integer) -> [Symbol?, Array[component]]?
 
+    private
+
+    def find_in_directive: (Integer pos, AST::Directives::t, Array[component]) -> bool
+
     def find_in_decl: (Integer pos, decl: AST::Declarations::t, array: Array[component]) -> bool
 
     def find_in_member: (Integer pos, member: AST::Members::t, array: Array[component]) -> bool

test/rbs/locator_test.rb

@@ -5,8 +5,8 @@ class RBS::LocatorTest < Test::Unit::TestCase
   include TestHelper
 
   def locator(src)
-    _, _, decls = Parser.parse_signature(src)
-    Locator.new(decls: decls)
+    buffer, dirs, decls = Parser.parse_signature(src)
+    Locator.new(buffer: buffer, dirs: dirs, decls: decls)
   end
 
   def test_find_class