Improve semantic token latency, highlight function and modified variables

expand.rkt

@@ -0,0 +1,97 @@
+#lang racket/base
+
+(provide read-and-expand)
+
+(require racket/match
+         racket/async-channel
+         syntax/modread
+         syntax/parse
+         drracket/check-syntax)
+
+(define *expander-ch* (make-async-channel))
+
+;; (-> input-port path collector% async-channel)
+;; async-channel receives (list (or syntax #f) (or expand-syntax #f))
+;; where `syntax` is the result of read-syntax,
+;; `expand-syntax` is the result of expand.
+;; They can be #f if error happens during processing.
+(define (read-and-expand in path collector)
+  (define ch (make-async-channel))
+  (async-channel-put *expander-ch* (list path in ch collector))
+  ch)
+
+(define (expander)
+  ;; cache namespace for continuous expand request for same uri
+  (define ns (make-base-namespace))
+  (define last-path #f)
+
+  (let loop ()
+    (match-define (list path in out-ch collector) (sync *expander-ch*))
+
+    (define result
+      (cond [(equal? last-path path)
+             (real-expand path in ns collector)]
+            [else
+             (set! ns (make-base-namespace))
+             (real-expand path in ns collector)]))
+
+    (set! last-path path)
+
+    (async-channel-put out-ch result)
+    (loop)))
+
+(define (real-expand path in ns collector)
+  (define-values (src-dir _1 _2) (split-path path))
+  (define-values (add-syntax done)
+    (make-traversal ns src-dir))
+
+  (parameterize ([current-load-relative-directory src-dir]
+                 [current-namespace ns]
+                 [current-annotations collector])
+    (define stx
+      (with-handlers ([(λ _ #t) (λ (exn) exn)])
+        (with-module-reading-parameterization
+          (lambda () (read-syntax path in)))))
+    (define expanded
+      (with-handlers ([(λ _ #t) (λ (exn) exn)])
+        (with-handlers ([(λ _ #t) (λ _ (expand stx))])
+          (expand (simplify-stx stx)))))
+
+    (when (not (exn? expanded))
+      (add-syntax expanded)
+      (done))
+
+    (list stx expanded)))
+
+(define _expand-th (thread expander))
+
+;; simplify syntax to optimize expand
+;; for example, use typed/racket/no-check language to avoid
+;; type check.
+(define (simplify-stx stx)
+  (define (apply-rules modname)
+    (match modname
+      [(or 'typed/racket/base
+           'typed/racket/base/deep
+           'typed/racket/base/shallow
+           'typed/racket/base/optional)
+       'typed/racket/base/no-check]
+      [(or 'typed/racket
+           'typed/racket/deep
+           'typed/racket/shallow
+           'typed/racket/optional)
+       'typed/racket/no-check]
+      [_ modname]))
+
+  (define (convert name-stx)
+    (define name-sym (syntax->datum name-stx))
+    (if (not (symbol? name-sym))
+        name-stx
+        (datum->syntax name-stx (apply-rules name-sym))))
+
+  (syntax-parse stx
+    [(module id mod-path form ...)
+     (define new-mod (convert #'mod-path))
+     #`(module id #,new-mod form ...)]
+    [stx #'stx]))
+

highlight.rkt

@@ -1,14 +1,14 @@
 #lang racket/base
 
-(require syntax/modread
-         drracket/check-syntax
+(require drracket/check-syntax
          syntax/parse
          "struct.rkt"
          racket/class
          racket/set
          racket/list
          racket/bool
-         racket/match)
+         racket/match
+         "expand.rkt")
 
 (provide collect-semantic-tokens)
 
@@ -44,31 +44,15 @@
   (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-styles)))
-    (set! token-list (append drracket-styles token-list)))
+  (match-define (list stx expanded) (sync (read-and-expand in path collector)))
+  (define drracket-styles (convert-drracket-color-styles (send collector get-styles)))
+
+  (define token-list
+    (append drracket-styles
+            (if (syntax? stx) (walk-stx stx) '())
+            (if (syntax? expanded) (walk-expanded-stx path expanded) '())))
 
   (let* ([tokens-no-false (filter-not false? token-list)]
          [tokens-no-out-bounds (filter (λ (t) (< -1 (Token-start t) (string-length code-str)))
@@ -92,6 +76,8 @@
     (match s
       [(Token start end 'drracket:check-syntax:lexically-bound)
        (Token start end 'variable)]
+      [(Token start end 'drracket:check-syntax:set!d)
+       (Token start end 'variable)]
       [_ #f])))
 
 ;; `tags` might contains multiple valid types.
@@ -125,12 +111,17 @@
 
 (define (walk-expanded-stx src stx)
   (syntax-parse stx
-    #:datum-literals (lambda define-values)
+    #:datum-literals (lambda define-values #%app)
     [(lambda (args ...) expr ...)
      (walk-expanded-stx src #'(expr ...))]
     [(define-values (fs) (lambda _ ...))
      (append (tags-of-stx-lst src #'(fs) 'function)
              (walk-expanded-stx src (drop (syntax-e stx) 2)))]
+    [(define-values (names ...) expr)
+     (walk-expanded-stx src #'expr)]
+    [(#%app proc args ...)
+     (append (tags-of-stx-lst src #'(proc) 'function)
+             (walk-expanded-stx src #'(args ...)))]
     [(any1 any* ...)
      (append (walk-expanded-stx src #'any1)
              (walk-expanded-stx src #'(any* ...)))]