From a45096b898f36b321a19ab328610bdd48786ad3c Mon Sep 17 00:00:00 2001
From: 6cdh <cd6cdh@gmail.com>
Date: Tue, 10 Sep 2024 00:20:56 +0800
Subject: [PATCH 1/6] feat(method): add initial semantic token support

---
 doc.rkt           |  39 +++++++++++-
 highlight.rkt     | 157 ++++++++++++++++++++++++++++++++++++++++++++++
 methods.rkt       |   8 +++
 struct.rkt        |  19 +++++-
 text-document.rkt |  11 +++-
 5 files changed, 231 insertions(+), 3 deletions(-)
 create mode 100644 highlight.rkt

diff --git a/doc.rkt b/doc.rkt
index ee1d548..bd1d1b8 100644
--- a/doc.rkt
+++ b/doc.rkt
@@ -9,6 +9,7 @@
          "path-util.rkt"
          "doc-trace.rkt"
          "struct.rkt"
+         "highlight.rkt"
          racket/match
          racket/class
          racket/set
@@ -315,6 +316,40 @@
                                   #:end (Pos #:line line #:char really-indent))
                    #:newText new-text)]))
 
+(define (token-type-encoding token)
+  (index-of *semantic-token-types* (SemanticToken-type token)))
+
+(define (token-modifier-encoding token)
+  (define indexes (indexes-where *semantic-token-modifiers*
+                                 (λ (m) (memq m (SemanticToken-modifiers token)))))
+  (for/sum ([index indexes])
+    (expt 2 index)))
+
+(define (token-encoding doc token prev-pos)
+  (define-values (line ch) (doc-line/ch doc (SemanticToken-start token)))
+  (define-values (prev-line prev-ch) (doc-line/ch doc prev-pos))
+  (define delta-line (- line prev-line))
+  (define delta-start
+    (if (= line prev-line)
+        (- ch prev-ch)
+        ch))
+  (define len (- (SemanticToken-end token) (SemanticToken-start token)))
+  (define type (token-type-encoding token))
+  (define modifier (token-modifier-encoding token))
+  (values delta-line delta-start len type modifier))
+
+(define (doc-full-tokens doc path)
+  (define tokens (collect-semantic-tokens (Doc-text doc) (uri->path path)))
+  (for/fold ([result '()]
+             [prev-pos 0]
+             #:result (let ()
+                        (flatten (reverse result))))
+            ([token tokens])
+    (define-values (delta-line delta-start len type modifier)
+      (token-encoding doc token prev-pos))
+    (values (cons (list delta-line delta-start len type modifier) result)
+            (SemanticToken-start token))))
+
 (provide Doc-trace
          new-doc
          doc-checked?
@@ -329,4 +364,6 @@
          doc-find-containing-paren
          doc-get-symbols
          get-definition-by-id
-         format!)
+         format!
+         doc-full-tokens)
+
diff --git a/highlight.rkt b/highlight.rkt
new file mode 100644
index 0000000..a45679e
--- /dev/null
+++ b/highlight.rkt
@@ -0,0 +1,157 @@
+#lang racket/base
+
+(require syntax/modread
+         drracket/check-syntax
+         syntax/parse
+         "struct.rkt"
+         racket/class
+         racket/set
+         racket/list
+         racket/bool
+         racket/match)
+
+(provide collect-semantic-tokens)
+
+(define collector%
+  (class (annotations-mixin object%)
+    (define styles '())
+
+    (super-new)
+
+    (define/override (syncheck:find-source-object stx)
+      #f)
+
+    (define/override (syncheck:color-range src start end style)
+      (when (< start end)
+        (set! styles (cons (list start end style) styles))))
+
+    (define/override (syncheck:add-definition-target src start finish id mods)
+      (when (< start finish)
+        (set! styles (cons (list start finish 'definition) styles))))
+
+    (define/public (get-color)
+      (set->list (list->set styles)))))
+
+; (-> lsp-editor% Path (Listof SemanticToken))
+(define (collect-semantic-tokens doc-text path)
+  (define code-str (send doc-text get-text))
+  (define in (open-input-string code-str))
+  (port-count-lines! in)
+  (define-values (path-dir _1 _2) (split-path path))
+
+  (define base-ns (make-base-namespace))
+
+  (define-values (add-syntax done)
+    (make-traversal base-ns #f))
+
+  (define token-list '())
+
+  (define collector (new collector%))
+  (with-handlers ([(λ (_) #t) (λ (_) #f)])
+    (parameterize ([current-load-relative-directory path-dir]
+                   [current-namespace base-ns]
+                   [current-annotations collector])
+      (define stx (with-module-reading-parameterization
+                    (lambda () (read-syntax path in))))
+      (set! token-list (append (walk-stx stx) token-list))
+
+      (define expanded (expand stx))
+      (set! token-list (append (walk-expanded-stx path expanded) token-list))
+      (add-syntax expanded)
+      (done))
+
+    (define drracket-styles (convert-drracket-color-styles (send collector get-color)))
+    (set! token-list (append drracket-styles token-list)))
+
+  (let* ([tokens-no-false (filter-not false? token-list)]
+         [tokens-no-out-bounds (filter (λ (t) (< -1 (first t) (string-length code-str)))
+                                       tokens-no-false)]
+         [tokens-in-order (sort tokens-no-out-bounds < #:key first)]
+         [same-loc-token-groups (group-by first tokens-in-order)]
+         [tokens-merge-types
+          (for/list ([group same-loc-token-groups])
+            (define fst (first group))
+            (list (first fst) (second fst) (map third group)))]
+         [result-tokens
+          (for*/list ([t tokens-merge-types]
+                      [type (in-value (select-type (third t)))]
+                      [modifiers (in-value (filter (λ (t) (memq t *semantic-token-modifiers*))
+                                                   (third t)))]
+                      #:when (not (false? type)))
+            (SemanticToken (first t) (second t) type modifiers))])
+    result-tokens))
+
+(define (convert-drracket-color-styles styles)
+  (for/list ([s styles])
+    (match s
+      [(list start end "drracket:check-syntax:lexically-bound")
+       (list start end 'variable)]
+      [_ #f])))
+
+(define (select-type types)
+  (define valid-types (filter (λ (t) (memq t *semantic-token-types*)) types))
+  (cond [(null? valid-types)
+         #f]
+        [(memq 'function valid-types)
+         'function]
+        [(memq 'variable valid-types)
+         'variable]
+        [else (first valid-types)]))
+
+(define (walk-stx stx)
+  (syntax-parse stx
+    #:datum-literals (#%module-begin)
+    [() (list)]
+    [(any1 any* ...)
+     (append (walk-stx #'any1)
+             (walk-stx #'(any* ...)))]
+    [#(any1 any* ...)
+     (append (walk-stx #'any1)
+             (walk-stx #'(any* ...)))]
+    [#%module-begin
+     (list)]
+    [atom (list (stx-typeof #'atom))]))
+
+(define (walk-expanded-stx src stx)
+  (syntax-parse stx
+    #:datum-literals (lambda define-values)
+    [(lambda (args ...) expr ...)
+     (walk-expanded-stx src #'(expr ...))]
+    [(define-values (fs) (lambda _ ...))
+     (append (stx-lst-typeof src #'(fs) 'function)
+             (walk-expanded-stx src (drop (syntax-e stx) 2)))]
+    [(any1 any* ...)
+     (append (walk-expanded-stx src #'any1)
+             (walk-expanded-stx src #'(any* ...)))]
+    [_ (list)]))
+
+(define (stx-lst-typeof src stx-lst type)
+  (define (in-current-file? stx)
+    (equal? src (syntax-source stx)))
+
+  (let* ([stx-lst (syntax-e stx-lst)]
+         [stx-lst-in-current-file (filter in-current-file? stx-lst)]
+         [type-lst (map (λ (stx) (stx-typeof stx type)) stx-lst-in-current-file)])
+    type-lst))
+
+(define (stx-typeof atom-stx [expect-type #f])
+  (define pos+1 (syntax-position atom-stx))
+  (define len (syntax-span atom-stx))
+  (if (or (not pos+1) (not len) (= len 0)
+          (not (syntax-original? atom-stx)))
+      #f
+      (let ([pos (sub1 pos+1)])
+        (list pos (+ pos len)
+              (if (false? expect-type)
+                  (get-type (syntax-e atom-stx))
+                  expect-type)))))
+
+(define (get-type atom)
+  (match atom
+    [(? number?) 'number]
+    [(? symbol?) 'symbol]
+    [(? string?) 'string]
+    [(? bytes?) 'string]
+    [(? regexp?) 'regexp]
+    [_ 'unknown]))
+
diff --git a/methods.rkt b/methods.rkt
index 7e87ee1..c637aa0 100644
--- a/methods.rkt
+++ b/methods.rkt
@@ -6,6 +6,7 @@
          "error-codes.rkt"
          "msg-io.rkt"
          "responses.rkt"
+         "struct.rkt"
          (prefix-in text-document/ "text-document.rkt"))
 
 ;; TextDocumentSynKind enumeration
@@ -88,6 +89,8 @@
        (text-document/range-formatting! id params)]
       ["textDocument/onTypeFormatting"
        (text-document/on-type-formatting! id params)]
+      ["textDocument/semanticTokens/full"
+       (text-document/full-semantic-tokens id params)]
       [_
        (eprintf "invalid request for method ~v\n" method)
        (define err (format "The method ~v was not found" method))
@@ -127,6 +130,10 @@
                                     (hash-table ['prepareSupport #t])])])
           (hasheq 'prepareProvider #t)]
          [_ #t]))
+     (define semantic-provider
+       (hasheq 'legend (hasheq 'tokenTypes (map symbol->string *semantic-token-types*)
+                               'tokenModifiers (map symbol->string *semantic-token-modifiers*))
+               'full #t))
      (define server-capabilities
        (hasheq 'textDocumentSync sync-options
                'hoverProvider #t
@@ -137,6 +144,7 @@
                'signatureHelpProvider (hasheq 'triggerCharacters (list " " ")" "]"))
                'inlayHintProvider #t
                'renameProvider renameProvider
+               'semanticTokensProvider semantic-provider
                'documentHighlightProvider #t
                'documentSymbolProvider #t
                'documentFormattingProvider #t
diff --git a/struct.rkt b/struct.rkt
index 7716fb7..bb6afd1 100644
--- a/struct.rkt
+++ b/struct.rkt
@@ -129,6 +129,20 @@
                             #:trim-final-newlines (hash-ref jsexpr 'trimFinalNewlines undef-object)
                             #:key (hash-ref jsexpr 'key undef-object))))
 
+(struct SemanticToken
+  (start end type modifiers)
+  #:transparent)
+
+(define *semantic-token-types*
+  '(variable
+     function
+     string
+     number
+     regexp))
+
+(define *semantic-token-modifiers*
+  '(definition))
+
 ;; usage:
 ;; (jsexpr? jsexpr) ;; #t
 ;; (match jsexpr
@@ -147,4 +161,7 @@
 (provide FormattingOptions
          FormattingOptions-tab-size
          FormattingOptions-trim-trailing-whitespace
-         as-FormattingOptions)
+         as-FormattingOptions
+         (struct-out SemanticToken)
+         *semantic-token-types*
+         *semantic-token-modifiers*)
diff --git a/text-document.rkt b/text-document.rkt
index 473138e..7dda7ba 100644
--- a/text-document.rkt
+++ b/text-document.rkt
@@ -478,6 +478,13 @@
     [_
      (error-response id INVALID-PARAMS "textDocument/onTypeFormatting failed")]))
 
+(define (full-semantic-tokens id params)
+  (match params
+    [(hash* ['textDocument (DocIdentifier #:uri uri)])
+     (define this-doc (hash-ref open-docs (string->symbol uri)))
+     (success-response id (hash 'data (doc-full-tokens this-doc uri)))]
+    [_ (error-response id INVALID-PARAMS "textDocument/semanticTokens/full failed")]))
+
 ;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;
 
 (provide
@@ -498,4 +505,6 @@
   [prepareRename (exact-nonnegative-integer? jsexpr? . -> . jsexpr?)]
   [formatting! (exact-nonnegative-integer? jsexpr? . -> . jsexpr?)]
   [range-formatting! (exact-nonnegative-integer? jsexpr? . -> . jsexpr?)]
-  [on-type-formatting! (exact-nonnegative-integer? jsexpr? . -> . jsexpr?)]))
+  [on-type-formatting! (exact-nonnegative-integer? jsexpr? . -> . jsexpr?)]
+  [full-semantic-tokens (exact-nonnegative-integer? jsexpr? . -> . jsexpr?)]))
+

From 55991664026cfc93b0e2c786d64355bf17630ad7 Mon Sep 17 00:00:00 2001
From: 6cdh <cd6cdh@gmail.com>
Date: Tue, 10 Sep 2024 00:54:59 +0800
Subject: [PATCH 2/6] feat(method): support range semantic token

---
 doc.rkt           | 10 +++++++---
 methods.rkt       |  5 ++++-
 text-document.rkt | 16 ++++++++++++++--
 3 files changed, 25 insertions(+), 6 deletions(-)

diff --git a/doc.rkt b/doc.rkt
index bd1d1b8..101c037 100644
--- a/doc.rkt
+++ b/doc.rkt
@@ -338,13 +338,17 @@
   (define modifier (token-modifier-encoding token))
   (values delta-line delta-start len type modifier))
 
-(define (doc-full-tokens doc path)
+(define (doc-range-tokens doc path pos-start pos-end)
   (define tokens (collect-semantic-tokens (Doc-text doc) (uri->path path)))
+  (define tokens-in-range
+    (filter-not (λ (tok) (or (<= (SemanticToken-end tok) pos-start)
+                             (>= (SemanticToken-start tok) pos-end)))
+                tokens))
   (for/fold ([result '()]
              [prev-pos 0]
              #:result (let ()
                         (flatten (reverse result))))
-            ([token tokens])
+            ([token tokens-in-range])
     (define-values (delta-line delta-start len type modifier)
       (token-encoding doc token prev-pos))
     (values (cons (list delta-line delta-start len type modifier) result)
@@ -365,5 +369,5 @@
          doc-get-symbols
          get-definition-by-id
          format!
-         doc-full-tokens)
+         doc-range-tokens)
 
diff --git a/methods.rkt b/methods.rkt
index c637aa0..dd93354 100644
--- a/methods.rkt
+++ b/methods.rkt
@@ -91,6 +91,8 @@
        (text-document/on-type-formatting! id params)]
       ["textDocument/semanticTokens/full"
        (text-document/full-semantic-tokens id params)]
+      ["textDocument/semanticTokens/range"
+       (text-document/range-semantic-tokens id params)]
       [_
        (eprintf "invalid request for method ~v\n" method)
        (define err (format "The method ~v was not found" method))
@@ -133,7 +135,8 @@
      (define semantic-provider
        (hasheq 'legend (hasheq 'tokenTypes (map symbol->string *semantic-token-types*)
                                'tokenModifiers (map symbol->string *semantic-token-modifiers*))
-               'full #t))
+               'full #t
+               'range #t))
      (define server-capabilities
        (hasheq 'textDocumentSync sync-options
                'hoverProvider #t
diff --git a/text-document.rkt b/text-document.rkt
index 7dda7ba..999096b 100644
--- a/text-document.rkt
+++ b/text-document.rkt
@@ -482,7 +482,18 @@
   (match params
     [(hash* ['textDocument (DocIdentifier #:uri uri)])
      (define this-doc (hash-ref open-docs (string->symbol uri)))
-     (success-response id (hash 'data (doc-full-tokens this-doc uri)))]
+     (success-response id (hash 'data (doc-range-tokens this-doc uri 0 (doc-endpos this-doc))))]
+    [_ (error-response id INVALID-PARAMS "textDocument/semanticTokens/full failed")]))
+
+(define (range-semantic-tokens id params)
+  (match params
+    [(hash* ['textDocument (DocIdentifier #:uri uri)]
+            ['range (Range #:start (Pos #:line st-ln #:char st-ch)
+                           #:end (Pos #:line ed-ln #:char ed-ch))])
+     (define this-doc (hash-ref open-docs (string->symbol uri)))
+     (define start-pos (doc-pos this-doc st-ln st-ch))
+     (define end-pos (doc-pos this-doc ed-ln ed-ch))
+     (success-response id (hash 'data (doc-range-tokens this-doc uri start-pos end-pos)))]
     [_ (error-response id INVALID-PARAMS "textDocument/semanticTokens/full failed")]))
 
 ;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;
@@ -506,5 +517,6 @@
   [formatting! (exact-nonnegative-integer? jsexpr? . -> . jsexpr?)]
   [range-formatting! (exact-nonnegative-integer? jsexpr? . -> . jsexpr?)]
   [on-type-formatting! (exact-nonnegative-integer? jsexpr? . -> . jsexpr?)]
-  [full-semantic-tokens (exact-nonnegative-integer? jsexpr? . -> . jsexpr?)]))
+  [full-semantic-tokens (exact-nonnegative-integer? jsexpr? . -> . jsexpr?)]
+  [range-semantic-tokens (exact-nonnegative-integer? jsexpr? . -> . jsexpr?)]))
 

From 30d4efa5ad3f5111ee1e17e04297b689211eb650 Mon Sep 17 00:00:00 2001
From: 6cdh <cd6cdh@gmail.com>
Date: Tue, 10 Sep 2024 01:01:01 +0800
Subject: [PATCH 3/6] fix: test

---
 tests/lifecycle/init_resp.json | 18 +++++++++++++++++-
 1 file changed, 17 insertions(+), 1 deletion(-)

diff --git a/tests/lifecycle/init_resp.json b/tests/lifecycle/init_resp.json
index f9b9909..c5c1a57 100644
--- a/tests/lifecycle/init_resp.json
+++ b/tests/lifecycle/init_resp.json
@@ -25,6 +25,22 @@
             "inlayHintProvider": true,
             "referencesProvider": true,
             "renameProvider": true,
+            "semanticTokensProvider": {
+                "full": true,
+                "range": true,
+                "legend": {
+                    "tokenModifiers": [
+                        "definition"
+                    ],
+                    "tokenTypes": [
+                        "variable",
+                        "function",
+                        "string",
+                        "number",
+                        "regexp"
+                    ]
+                }
+            },
             "signatureHelpProvider": {
                 "triggerCharacters": [
                     " ",
@@ -40,4 +56,4 @@
             }
         }
     }
-}
\ No newline at end of file
+}

From 6ead0ae3cd0f7f0181264aefb647ff284f9ea226 Mon Sep 17 00:00:00 2001
From: 6cdh <cd6cdh@gmail.com>
Date: Tue, 10 Sep 2024 01:34:34 +0800
Subject: [PATCH 4/6] fix(ci): use hash-table instead of hash*

---
 text-document.rkt | 6 +++---
 1 file changed, 3 insertions(+), 3 deletions(-)

diff --git a/text-document.rkt b/text-document.rkt
index 999096b..9890579 100644
--- a/text-document.rkt
+++ b/text-document.rkt
@@ -480,15 +480,15 @@
 
 (define (full-semantic-tokens id params)
   (match params
-    [(hash* ['textDocument (DocIdentifier #:uri uri)])
+    [(hash-table ['textDocument (DocIdentifier #:uri uri)])
      (define this-doc (hash-ref open-docs (string->symbol uri)))
      (success-response id (hash 'data (doc-range-tokens this-doc uri 0 (doc-endpos this-doc))))]
     [_ (error-response id INVALID-PARAMS "textDocument/semanticTokens/full failed")]))
 
 (define (range-semantic-tokens id params)
   (match params
-    [(hash* ['textDocument (DocIdentifier #:uri uri)]
-            ['range (Range #:start (Pos #:line st-ln #:char st-ch)
+    [(hash-table ['textDocument (DocIdentifier #:uri uri)]
+                 ['range (Range #:start (Pos #:line st-ln #:char st-ch)
                            #:end (Pos #:line ed-ln #:char ed-ch))])
      (define this-doc (hash-ref open-docs (string->symbol uri)))
      (define start-pos (doc-pos this-doc st-ln st-ch))

From e410cd34b8aa80f466e8d59014acfc114938f38d Mon Sep 17 00:00:00 2001
From: 6cdh <cd6cdh@gmail.com>
Date: Tue, 10 Sep 2024 20:53:00 +0800
Subject: [PATCH 5/6] style: improve code, add some comment

---
 doc.rkt                        | 31 ++++++++++++++-
 highlight.rkt                  | 69 ++++++++++++++++++++--------------
 struct.rkt                     |  7 ++++
 tests/lifecycle/init_resp.json | 14 +------
 tests/lifecycle/test-main.rkt  | 11 +++++-
 text-document.rkt              |  2 +-
 6 files changed, 87 insertions(+), 47 deletions(-)

diff --git a/doc.rkt b/doc.rkt
index 101c037..c2332f3 100644
--- a/doc.rkt
+++ b/doc.rkt
@@ -322,9 +322,33 @@
 (define (token-modifier-encoding token)
   (define indexes (indexes-where *semantic-token-modifiers*
                                  (λ (m) (memq m (SemanticToken-modifiers token)))))
+  ;; build a bit flag of the modifiers of `token`.
+  ;;
+  ;; equivalent to C family pseudocode
+  ;;
+  ;; uint32_t flag = 0
+  ;; for index in indexes:
+  ;;   flag = flag | (1 << index)
+  ;; return flag
+  ;;
+  ;; But the integer bit width is ignored here, because
+  ;; the *semantic-token-modifiers* is very small.
   (for/sum ([index indexes])
     (expt 2 index)))
 
+;; encode `token` using relative encoding
+;;
+;; each token is encoded as five integers (copied from lsp specificatioin 3.17):
+;; * deltaLine: token line number, relative to the start of the previous token
+;; * deltaStart: token start character, relative to the start of the previous token 
+;;               (relative to 0 or the previous token’s start if they are on the same line)
+;; * length: the length of the token.
+;; * tokenType: will be looked up in SemanticTokensLegend.tokenTypes.
+;;              We currently ask that tokenType < 65536.
+;; * tokenModifiers: each set bit will be looked up in SemanticTokensLegend.tokenModifiers
+;;
+;; for the first token, its previous token is defined as a zero length fake token which
+;; has line number 0 and character position 0.
 (define (token-encoding doc token prev-pos)
   (define-values (line ch) (doc-line/ch doc (SemanticToken-start token)))
   (define-values (prev-line prev-ch) (doc-line/ch doc prev-pos))
@@ -338,6 +362,10 @@
   (define modifier (token-modifier-encoding token))
   (values delta-line delta-start len type modifier))
 
+;; get the tokens whose range are contained in interval [pos-start, pos-end)
+;; the tokens whose range intersects the given range is included.
+;; the previous token of the first token in the result is defined as a zero length fake token which
+;; has line number 0 and character position 0.
 (define (doc-range-tokens doc path pos-start pos-end)
   (define tokens (collect-semantic-tokens (Doc-text doc) (uri->path path)))
   (define tokens-in-range
@@ -346,8 +374,7 @@
                 tokens))
   (for/fold ([result '()]
              [prev-pos 0]
-             #:result (let ()
-                        (flatten (reverse result))))
+             #:result (flatten (reverse result)))
             ([token tokens-in-range])
     (define-values (delta-line delta-start len type modifier)
       (token-encoding doc token prev-pos))
diff --git a/highlight.rkt b/highlight.rkt
index a45679e..8067e65 100644
--- a/highlight.rkt
+++ b/highlight.rkt
@@ -12,6 +12,13 @@
 
 (provide collect-semantic-tokens)
 
+;; A temporary structure to hold tokens
+;; `tag` is symbol that is a tag associated with this token.
+;; An identifier may correspond multiple tokens. They will be merged, then converted into
+;; lsp semantic token types and modifiers.
+(struct Token
+  (start end tag))
+
 (define collector%
   (class (annotations-mixin object%)
     (define styles '())
@@ -23,13 +30,13 @@
 
     (define/override (syncheck:color-range src start end style)
       (when (< start end)
-        (set! styles (cons (list start end style) styles))))
+        (set! styles (cons (Token start end (string->symbol style)) styles))))
 
     (define/override (syncheck:add-definition-target src start finish id mods)
       (when (< start finish)
-        (set! styles (cons (list start finish 'definition) styles))))
+        (set! styles (cons (Token start finish 'definition) styles))))
 
-    (define/public (get-color)
+    (define/public (get-styles)
       (set->list (list->set styles)))))
 
 ; (-> lsp-editor% Path (Listof SemanticToken))
@@ -60,23 +67,22 @@
       (add-syntax expanded)
       (done))
 
-    (define drracket-styles (convert-drracket-color-styles (send collector get-color)))
+    (define drracket-styles (convert-drracket-color-styles (send collector get-styles)))
     (set! token-list (append drracket-styles token-list)))
 
   (let* ([tokens-no-false (filter-not false? token-list)]
-         [tokens-no-out-bounds (filter (λ (t) (< -1 (first t) (string-length code-str)))
+         [tokens-no-out-bounds (filter (λ (t) (< -1 (Token-start t) (string-length code-str)))
                                        tokens-no-false)]
-         [tokens-in-order (sort tokens-no-out-bounds < #:key first)]
-         [same-loc-token-groups (group-by first tokens-in-order)]
-         [tokens-merge-types
-          (for/list ([group same-loc-token-groups])
-            (define fst (first group))
-            (list (first fst) (second fst) (map third group)))]
+         [tokens-in-order (sort tokens-no-out-bounds < #:key Token-start)]
+         [same-ident-token-groups (group-by Token-start tokens-in-order)]
+         [tokens-with-merged-tags
+          (for/list ([token-group same-ident-token-groups])
+            (define tok (first token-group))
+            (list (Token-start tok) (Token-end tok) (map Token-tag token-group)))]
          [result-tokens
-          (for*/list ([t tokens-merge-types]
+          (for*/list ([t tokens-with-merged-tags]
                       [type (in-value (select-type (third t)))]
-                      [modifiers (in-value (filter (λ (t) (memq t *semantic-token-modifiers*))
-                                                   (third t)))]
+                      [modifiers (in-value (get-valid-modifiers (third t)))]
                       #:when (not (false? type)))
             (SemanticToken (first t) (second t) type modifiers))])
     result-tokens))
@@ -84,12 +90,14 @@
 (define (convert-drracket-color-styles styles)
   (for/list ([s styles])
     (match s
-      [(list start end "drracket:check-syntax:lexically-bound")
-       (list start end 'variable)]
+      [(Token start end 'drracket:check-syntax:lexically-bound)
+       (Token start end 'variable)]
       [_ #f])))
 
-(define (select-type types)
-  (define valid-types (filter (λ (t) (memq t *semantic-token-types*)) types))
+;; `tags` might contains multiple valid types.
+;; This function selects a proper type based on some rules.
+(define (select-type tags)
+  (define valid-types (filter (λ (t) (memq t *semantic-token-types*)) tags))
   (cond [(null? valid-types)
          #f]
         [(memq 'function valid-types)
@@ -98,6 +106,9 @@
          'variable]
         [else (first valid-types)]))
 
+(define (get-valid-modifiers tags)
+  (filter (λ (t) (memq t *semantic-token-modifiers*)) tags))
+
 (define (walk-stx stx)
   (syntax-parse stx
     #:datum-literals (#%module-begin)
@@ -110,7 +121,7 @@
              (walk-stx #'(any* ...)))]
     [#%module-begin
      (list)]
-    [atom (list (stx-typeof #'atom))]))
+    [atom (list (tag-of-atom-stx #'atom))]))
 
 (define (walk-expanded-stx src stx)
   (syntax-parse stx
@@ -118,35 +129,35 @@
     [(lambda (args ...) expr ...)
      (walk-expanded-stx src #'(expr ...))]
     [(define-values (fs) (lambda _ ...))
-     (append (stx-lst-typeof src #'(fs) 'function)
+     (append (tags-of-stx-lst src #'(fs) 'function)
              (walk-expanded-stx src (drop (syntax-e stx) 2)))]
     [(any1 any* ...)
      (append (walk-expanded-stx src #'any1)
              (walk-expanded-stx src #'(any* ...)))]
     [_ (list)]))
 
-(define (stx-lst-typeof src stx-lst type)
+(define (tags-of-stx-lst src stx-lst tag)
   (define (in-current-file? stx)
     (equal? src (syntax-source stx)))
 
   (let* ([stx-lst (syntax-e stx-lst)]
          [stx-lst-in-current-file (filter in-current-file? stx-lst)]
-         [type-lst (map (λ (stx) (stx-typeof stx type)) stx-lst-in-current-file)])
-    type-lst))
+         [tag-lst (map (λ (stx) (tag-of-atom-stx stx tag)) stx-lst-in-current-file)])
+    tag-lst))
 
-(define (stx-typeof atom-stx [expect-type #f])
+(define (tag-of-atom-stx atom-stx [expect-tag #f])
   (define pos+1 (syntax-position atom-stx))
   (define len (syntax-span atom-stx))
   (if (or (not pos+1) (not len) (= len 0)
           (not (syntax-original? atom-stx)))
       #f
       (let ([pos (sub1 pos+1)])
-        (list pos (+ pos len)
-              (if (false? expect-type)
-                  (get-type (syntax-e atom-stx))
-                  expect-type)))))
+        (Token pos (+ pos len)
+               (if (false? expect-tag)
+                   (get-atom-tag (syntax-e atom-stx))
+                   expect-tag)))))
 
-(define (get-type atom)
+(define (get-atom-tag atom)
   (match atom
     [(? number?) 'number]
     [(? symbol?) 'symbol]
diff --git a/struct.rkt b/struct.rkt
index bb6afd1..ac9abfc 100644
--- a/struct.rkt
+++ b/struct.rkt
@@ -133,6 +133,12 @@
   (start end type modifiers)
   #:transparent)
 
+;; The order of this list is irrelevant.
+;; The client receives this list from server ability declaration during
+;; initialize handshake then use it to decode server semantic tokens messages.
+;; Different order produces different encoding results of semantic tokens,
+;; but does not affect client and server behavior.
+;; To change the order, simply change it here, don't need to change other code.
 (define *semantic-token-types*
   '(variable
      function
@@ -140,6 +146,7 @@
      number
      regexp))
 
+;; The order of this list is irrelevant, similar to *semantic-token-types*.
 (define *semantic-token-modifiers*
   '(definition))
 
diff --git a/tests/lifecycle/init_resp.json b/tests/lifecycle/init_resp.json
index c5c1a57..3f7a8e9 100644
--- a/tests/lifecycle/init_resp.json
+++ b/tests/lifecycle/init_resp.json
@@ -27,19 +27,7 @@
             "renameProvider": true,
             "semanticTokensProvider": {
                 "full": true,
-                "range": true,
-                "legend": {
-                    "tokenModifiers": [
-                        "definition"
-                    ],
-                    "tokenTypes": [
-                        "variable",
-                        "function",
-                        "string",
-                        "number",
-                        "regexp"
-                    ]
-                }
+                "range": true
             },
             "signatureHelpProvider": {
                 "triggerCharacters": [
diff --git a/tests/lifecycle/test-main.rkt b/tests/lifecycle/test-main.rkt
index a8340aa..b4d8d59 100644
--- a/tests/lifecycle/test-main.rkt
+++ b/tests/lifecycle/test-main.rkt
@@ -2,7 +2,9 @@
 (require chk
          json
          racket/os
-         "../../msg-io.rkt")
+         "../../msg-io.rkt"
+         "../../json-util.rkt"
+         "../../struct.rkt")
 
 (define init-req
   (hasheq 'jsonrpc "2.0"
@@ -35,7 +37,11 @@
   ;; Initialize request
   (display-message/flush init-req stdin)
   (let ([resp (read-message stdout)])
-    (chk #:= resp (read-json (open-input-file "init_resp.json"))))
+    (define expected-json (read-json (open-input-file "init_resp.json")))
+    (define json (jsexpr-set expected-json '(result capabilities semanticTokensProvider legend)
+                             (hasheq 'tokenModifiers (map symbol->string *semantic-token-modifiers*)
+                                     'tokenTypes (map symbol->string *semantic-token-types*))))
+    (chk #:= resp json))
 
   ;; Shutdown request
   (display-message/flush shutdown-req stdin)
@@ -49,3 +55,4 @@
   (subprocess-wait sp)
   (define st (subprocess-status sp))
   (chk (zero? st)))
+
diff --git a/text-document.rkt b/text-document.rkt
index 9890579..8811181 100644
--- a/text-document.rkt
+++ b/text-document.rkt
@@ -494,7 +494,7 @@
      (define start-pos (doc-pos this-doc st-ln st-ch))
      (define end-pos (doc-pos this-doc ed-ln ed-ch))
      (success-response id (hash 'data (doc-range-tokens this-doc uri start-pos end-pos)))]
-    [_ (error-response id INVALID-PARAMS "textDocument/semanticTokens/full failed")]))
+    [_ (error-response id INVALID-PARAMS "textDocument/semanticTokens/range failed")]))
 
 ;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;
 

From 883957b172068b5f7bb16dc1837accaef27a27e2 Mon Sep 17 00:00:00 2001
From: 6cdh <cd6cdh@gmail.com>
Date: Sun, 15 Sep 2024 01:01:35 +0800
Subject: [PATCH 6/6] feat(method): make semantic token method runs
 concurrently

---
 doc.rkt           | 15 ++++++++-------
 methods.rkt       | 10 +++++++++-
 text-document.rkt | 17 ++++++++++++-----
 3 files changed, 29 insertions(+), 13 deletions(-)

diff --git a/doc.rkt b/doc.rkt
index c2332f3..d3e47d4 100644
--- a/doc.rkt
+++ b/doc.rkt
@@ -349,9 +349,9 @@
 ;;
 ;; for the first token, its previous token is defined as a zero length fake token which
 ;; has line number 0 and character position 0.
-(define (token-encoding doc token prev-pos)
-  (define-values (line ch) (doc-line/ch doc (SemanticToken-start token)))
-  (define-values (prev-line prev-ch) (doc-line/ch doc prev-pos))
+(define (token-encoding editor token prev-pos)
+  (match-define (list line ch) (send editor pos->line/char (SemanticToken-start token)))
+  (match-define (list prev-line prev-ch) (send editor pos->line/char prev-pos))
   (define delta-line (- line prev-line))
   (define delta-start
     (if (= line prev-line)
@@ -366,8 +366,8 @@
 ;; the tokens whose range intersects the given range is included.
 ;; the previous token of the first token in the result is defined as a zero length fake token which
 ;; has line number 0 and character position 0.
-(define (doc-range-tokens doc path pos-start pos-end)
-  (define tokens (collect-semantic-tokens (Doc-text doc) (uri->path path)))
+(define (doc-range-tokens editor uri pos-start pos-end)
+  (define tokens (collect-semantic-tokens editor (uri->path uri)))
   (define tokens-in-range
     (filter-not (λ (tok) (or (<= (SemanticToken-end tok) pos-start)
                              (>= (SemanticToken-start tok) pos-end)))
@@ -377,11 +377,12 @@
              #:result (flatten (reverse result)))
             ([token tokens-in-range])
     (define-values (delta-line delta-start len type modifier)
-      (token-encoding doc token prev-pos))
+      (token-encoding editor token prev-pos))
     (values (cons (list delta-line delta-start len type modifier) result)
             (SemanticToken-start token))))
 
-(provide Doc-trace
+(provide Doc-text
+         Doc-trace
          new-doc
          doc-checked?
          doc-update!
diff --git a/methods.rkt b/methods.rkt
index dd93354..713a021 100644
--- a/methods.rkt
+++ b/methods.rkt
@@ -31,7 +31,15 @@
                  ['method (? string? method)])
      (define params (hash-ref msg 'params hasheq))
      (define response (process-request id method params))
-     (display-message/flush response)]
+     ;; the result can be a response or a procedure which returns
+     ;; a response. If it's a procedure, then it's expected to run
+     ;; concurrently.
+     (thread (λ ()
+               (display-message/flush
+                 (if (procedure? response)
+                     (response)
+                     response))))
+     (void)]
     ;; Notification
     [(hash-table ['method (? string? method)])
      (define params (hash-ref msg 'params hasheq))
diff --git a/text-document.rkt b/text-document.rkt
index 8811181..7a4530b 100644
--- a/text-document.rkt
+++ b/text-document.rkt
@@ -482,20 +482,27 @@
   (match params
     [(hash-table ['textDocument (DocIdentifier #:uri uri)])
      (define this-doc (hash-ref open-docs (string->symbol uri)))
-     (success-response id (hash 'data (doc-range-tokens this-doc uri 0 (doc-endpos this-doc))))]
+     (semantic-tokens uri id 0 (doc-endpos this-doc))]
     [_ (error-response id INVALID-PARAMS "textDocument/semanticTokens/full failed")]))
 
 (define (range-semantic-tokens id params)
   (match params
     [(hash-table ['textDocument (DocIdentifier #:uri uri)]
                  ['range (Range #:start (Pos #:line st-ln #:char st-ch)
-                           #:end (Pos #:line ed-ln #:char ed-ch))])
+                                #:end (Pos #:line ed-ln #:char ed-ch))])
      (define this-doc (hash-ref open-docs (string->symbol uri)))
      (define start-pos (doc-pos this-doc st-ln st-ch))
      (define end-pos (doc-pos this-doc ed-ln ed-ch))
-     (success-response id (hash 'data (doc-range-tokens this-doc uri start-pos end-pos)))]
+     (semantic-tokens uri id start-pos end-pos)]
     [_ (error-response id INVALID-PARAMS "textDocument/semanticTokens/range failed")]))
 
+(define (semantic-tokens uri id start-pos end-pos)
+  (define this-doc (hash-ref open-docs (string->symbol uri)))
+  (define new-editor (send (Doc-text this-doc) copy))
+  (λ ()
+    (define tokens (doc-range-tokens new-editor uri start-pos end-pos))
+    (success-response id (hash 'data tokens))))
+
 ;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;
 
 (provide
@@ -517,6 +524,6 @@
   [formatting! (exact-nonnegative-integer? jsexpr? . -> . jsexpr?)]
   [range-formatting! (exact-nonnegative-integer? jsexpr? . -> . jsexpr?)]
   [on-type-formatting! (exact-nonnegative-integer? jsexpr? . -> . jsexpr?)]
-  [full-semantic-tokens (exact-nonnegative-integer? jsexpr? . -> . jsexpr?)]
-  [range-semantic-tokens (exact-nonnegative-integer? jsexpr? . -> . jsexpr?)]))
+  [full-semantic-tokens (exact-nonnegative-integer? jsexpr? . -> . (or/c jsexpr? (-> jsexpr?)))]
+  [range-semantic-tokens (exact-nonnegative-integer? jsexpr? . -> . (or/c jsexpr? (-> jsexpr?)))]))