From e650df98c4ad02c9cdc791fb21c4b621f4635c9f Mon Sep 17 00:00:00 2001
From: Vinicius Stock <vinistock@users.noreply.github.com>
Date: Thu, 15 Aug 2024 16:15:11 -0400
Subject: [PATCH] Create RBSDocument (#2437)

---
 lib/ruby_lsp/document.rb     |  1 +
 lib/ruby_lsp/internal.rb     |  1 +
 lib/ruby_lsp/rbs_document.rb | 41 ++++++++++++++++++++++++++++++++++++
 lib/ruby_lsp/server.rb       |  3 +++
 lib/ruby_lsp/store.rb        | 10 ++++++---
 test/rbs_document_test.rb    | 35 ++++++++++++++++++++++++++++++
 6 files changed, 88 insertions(+), 3 deletions(-)
 create mode 100644 lib/ruby_lsp/rbs_document.rb
 create mode 100644 test/rbs_document_test.rb

diff --git a/lib/ruby_lsp/document.rb b/lib/ruby_lsp/document.rb
index 840bb9c6c..f6c503eb6 100644
--- a/lib/ruby_lsp/document.rb
+++ b/lib/ruby_lsp/document.rb
@@ -7,6 +7,7 @@ class LanguageId < T::Enum
       enums do
         Ruby = new("ruby")
         ERB = new("erb")
+        RBS = new("rbs")
       end
     end
 
diff --git a/lib/ruby_lsp/internal.rb b/lib/ruby_lsp/internal.rb
index 05fd4c91b..0f5ded017 100644
--- a/lib/ruby_lsp/internal.rb
+++ b/lib/ruby_lsp/internal.rb
@@ -36,6 +36,7 @@
 require "ruby_lsp/document"
 require "ruby_lsp/ruby_document"
 require "ruby_lsp/erb_document"
+require "ruby_lsp/rbs_document"
 require "ruby_lsp/store"
 require "ruby_lsp/addon"
 require "ruby_lsp/requests/support/rubocop_runner"
diff --git a/lib/ruby_lsp/rbs_document.rb b/lib/ruby_lsp/rbs_document.rb
new file mode 100644
index 000000000..92a8465b3
--- /dev/null
+++ b/lib/ruby_lsp/rbs_document.rb
@@ -0,0 +1,41 @@
+# typed: strict
+# frozen_string_literal: true
+
+module RubyLsp
+  class RBSDocument < Document
+    extend T::Sig
+    extend T::Generic
+
+    ParseResultType = type_member { { fixed: T::Array[RBS::AST::Declarations::Base] } }
+
+    sig { params(source: String, version: Integer, uri: URI::Generic, encoding: Encoding).void }
+    def initialize(source:, version:, uri:, encoding: Encoding::UTF_8)
+      @syntax_error = T.let(false, T::Boolean)
+      super
+    end
+
+    sig { override.returns(ParseResultType) }
+    def parse
+      return @parse_result unless @needs_parsing
+
+      @needs_parsing = false
+
+      _, _, declarations = RBS::Parser.parse_signature(@source)
+      @syntax_error = false
+      @parse_result = declarations
+    rescue RBS::ParsingError
+      @syntax_error = true
+      @parse_result
+    end
+
+    sig { override.returns(T::Boolean) }
+    def syntax_error?
+      @syntax_error
+    end
+
+    sig { override.returns(LanguageId) }
+    def language_id
+      LanguageId::RBS
+    end
+  end
+end
diff --git a/lib/ruby_lsp/server.rb b/lib/ruby_lsp/server.rb
index 9db8a6b4c..8ffb2b612 100644
--- a/lib/ruby_lsp/server.rb
+++ b/lib/ruby_lsp/server.rb
@@ -298,9 +298,12 @@ def text_document_did_open(message)
         language_id = case text_document[:languageId]
         when "erb", "eruby"
           Document::LanguageId::ERB
+        when "rbs"
+          Document::LanguageId::RBS
         else
           Document::LanguageId::Ruby
         end
+
         @store.set(
           uri: text_document[:uri],
           source: text_document[:text],
diff --git a/lib/ruby_lsp/store.rb b/lib/ruby_lsp/store.rb
index 0e1a5f2d5..80ac596a2 100644
--- a/lib/ruby_lsp/store.rb
+++ b/lib/ruby_lsp/store.rb
@@ -44,8 +44,11 @@ def get(uri)
       raise NonExistingDocumentError, uri.to_s unless path
 
       ext = File.extname(path)
-      language_id = if ext == ".erb" || ext == ".rhtml"
+      language_id = case ext
+      when ".erb", ".rhtml"
         Document::LanguageId::ERB
+      when ".rbs"
+        Document::LanguageId::RBS
       else
         Document::LanguageId::Ruby
       end
@@ -66,13 +69,14 @@ def get(uri)
       ).void
     end
     def set(uri:, source:, version:, language_id:, encoding: Encoding::UTF_8)
-      document = case language_id
+      @state[uri.to_s] = case language_id
       when Document::LanguageId::ERB
         ERBDocument.new(source: source, version: version, uri: uri, encoding: encoding)
+      when Document::LanguageId::RBS
+        RBSDocument.new(source: source, version: version, uri: uri, encoding: encoding)
       else
         RubyDocument.new(source: source, version: version, uri: uri, encoding: encoding)
       end
-      @state[uri.to_s] = document
     end
 
     sig { params(uri: URI::Generic, edits: T::Array[T::Hash[Symbol, T.untyped]], version: Integer).void }
diff --git a/test/rbs_document_test.rb b/test/rbs_document_test.rb
new file mode 100644
index 000000000..2206b31ac
--- /dev/null
+++ b/test/rbs_document_test.rb
@@ -0,0 +1,35 @@
+# typed: true
+# frozen_string_literal: true
+
+require "test_helper"
+
+class RBSDocumentTest < Minitest::Test
+  def test_parse_result_is_array_of_declarations
+    document = RubyLsp::RBSDocument.new(source: <<~RBS, version: 1, uri: URI("file:///foo.rbs"))
+      class Foo
+        def bar: () -> void
+      end
+    RBS
+
+    refute_predicate(document, :syntax_error?)
+    assert_equal(:Foo, T.cast(document.parse_result[0], RBS::AST::Declarations::Class).name.name)
+  end
+
+  def test_parsing_remembers_syntax_errors
+    # The syntax error is that `-` should be `->`
+    document = RubyLsp::RBSDocument.new(source: +<<~RBS, version: 1, uri: URI("file:///foo.rbs"))
+      class Foo
+        def bar: () - void
+      end
+    RBS
+
+    assert_predicate(document, :syntax_error?)
+
+    document.push_edits(
+      [{ range: { start: { line: 1, character: 15 }, end: { line: 1, character: 15 } }, text: ">" }],
+      version: 2,
+    )
+    document.parse
+    refute_predicate(document, :syntax_error?)
+  end
+end