racket semantic tokens

pt.rkt

#lang racket

(require syntax/modread
         racket/cmdline
         drracket/check-syntax
         syntax/parse)

(collect-garbage)

(define collector%
  (class (annotations-mixin object%)
    (init-field *src)
    (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)))))

(define-syntax-rule (with x expr ...)
  (let* ([x expr] ...)
    x))

(define (run path)
  (define src path)
  (define code (file->string path))
  (define in (open-input-file path))
  (port-count-lines! in)
  (define-values (src-dir _1 _2) (split-path src))

  (define base-ns (make-base-namespace))

  (define-values (add-syntax done)
    (make-traversal base-ns #f))

  (define semantic-infos '())

  (define collector (new collector% [*src src]))
  (parameterize ([current-load-relative-directory src-dir]
                 [current-namespace base-ns]
                 [current-annotations collector])
    (define stx (with-module-reading-parameterization
                  (lambda () (read-syntax src in))))
    (set! semantic-infos (append (walk-stx stx) semantic-infos))
    (define expanded (expand stx))
    (set! semantic-infos (append (walk-stx-expand src expanded) semantic-infos))
    (add-syntax expanded)
    (done))

  (set! semantic-infos (append (send collector get-color) semantic-infos))
  (with si
        (sort semantic-infos < #:key first)
        (group-by first si)
        (map (λ (g) (list (substring code
                                     (first (first g))
                                     (second (first g)))
                          (map third g)))
             si)))

(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 (add-stx #'atom)]))

(define (walk-stx-expand src stx)
  (syntax-parse stx
    #:datum-literals (lambda define-values)
    [(lambda (args ...) expr ...)
     (append (add-stx-lst-checked src #'(args ...) 'parameter)
             (walk-stx-expand src #'(expr ...)))]
    [(define-values (fs) (lambda _ ...))
     (append (add-stx-lst-checked src #'(fs) 'function)
             (walk-stx-expand src (drop (syntax-e stx) 2)))]
    [(any1 any* ...)
     (append (walk-stx-expand src #'any1)
             (walk-stx-expand src #'(any* ...)))]
    [_ (list)]))

(define (add-stx-lst-checked src stx-lst type)
  (with lst
        (syntax-e stx-lst)
        (filter (λ (s) (equal? src (syntax-source s))) lst)
        (filter (λ (s) (syntax-original? s)) lst)
        (map (λ (s) (add-stx-type s type)) lst)
        (apply append lst)))

(define (add-stx-type stx type)
  (define pos+1 (syntax-position stx))
  (define len (syntax-span stx))
  (if (or (not pos+1) (not len) (= len 0))
      (list)
      (let ([pos (sub1 pos+1)])
        (list (list pos (+ pos len)
                    (if (eq? type 'infer)
                        (typeof (syntax-e stx))
                        type))))))

(define (add-stx-checked stx type)
  (if (syntax-original? stx)
      (add-stx-type stx type)
      (list)))

(define (add-stx stx)
  (add-stx-checked stx 'infer))

(define (typeof datum)
  (match datum
    [(? boolean?) 'boolean]
    [(? number?) 'number]
    [(? symbol?) 'symbol]
    [(? keyword?) 'keyword]
    [(? string?) 'string]
    [(? bytes?) 'string]
    [(? regexp?) 'regexp]
    [_ 'UNKNOWN]))

(define filename (command-line #:args (filename) filename))
(define src (path->complete-path (string->path filename)))

(time (let ()
        (define infos (run src))
        (for ([info infos])
          (displayln (format "~a\t\t~a" (first info) (second info))))))