syntax-traversal.rkt

#lang racket/base


(require racket/contract/base)


(provide atom
         expression-directly-enclosing
         syntax-find-first
         syntax-search
         syntax-search-everything
         syntax-traverse)


(require (for-syntax racket/base
                     resyntax/private/more-syntax-parse-classes)
         racket/match
         racket/sequence
         racket/stream
         syntax/parse
         syntax/parse/define)


(module+ test
  (require racket/syntax-srcloc
           rackunit
           (submod "..")))


;@----------------------------------------------------------------------------------------------------


(define-syntax-class atom
  (pattern (~or :id :number :str :boolean :regexp :keyword)))


(define-syntax-class (expression-directly-enclosing id)
  (pattern (part ...)
    #:when (for/or ([part-stx (in-list (attribute part))]
                    #:when (identifier? part-stx))
             (free-identifier=? id part-stx)))
  (pattern (part ...+ . tail-part)
    #:when (for/or ([part-stx (in-list (cons #'tail-part (attribute part)))]
                    #:when (identifier? part-stx))
             (free-identifier=? id part-stx))))


(begin-for-syntax
  (define-syntax-class syntax-search-clause
    #:attributes (syntax-pattern [directive 1] output-stream)
    (pattern [syntax-pattern directive:syntax-parse-pattern-directive ... body:expr ... last-body]
      #:declare last-body (expr/c #'stream?)
      #:with output-stream #'(stream-lazy (let () body ... last-body.c)))
    (pattern [syntax-pattern directive:syntax-parse-pattern-directive ...]
      #:with output-stream #'(stream this-syntax))))


(define-syntax-parse-rule
  (syntax-search stx-expr
    (~optional (~seq #:skip-root? skip-root?) #:defaults ([skip-root? #'#false]))
    option:syntax-parse-option ...
    clause:syntax-search-clause ...)
  #:declare stx-expr (expr/c #'syntax?)
  (let ([skip-root-id skip-root?])
    (define-syntax-class search-case
      #:attributes (output-stream)
      (~@ . option) ...
      (pattern clause.syntax-pattern (~@ . clause.directive) ...
        #:attr output-stream clause.output-stream)
      ...)
    (let loop ([stx stx-expr.c] [root? #true])
      (stream-lazy
       (syntax-parse stx
         [child
          #:when (not (and skip-root-id root?))
          #:with (~var matched search-case) (attribute child)
          (attribute matched.output-stream)]
         [(part (... ...))
          #:cut
          (apply stream-append
                 (for/list ([part-stx (in-list (attribute part))])
                   (loop part-stx #false)))]
         [(part (... ...+) . tail-part)
          #:cut
          (stream-append (apply stream-append
                                (for/list ([part-stx (in-list (attribute part))])
                                  (loop part-stx #false)))
                         (loop #'tail-part #false))]
         [_ (stream)])))))


;; Returns stx and every syntax object nested within it, in depth-first preorder. This traversal is
;; eager, even though it's exposed as a stream, because building it lazily costs far more than the
;; list it saves allocating: every caller consumes nearly all of the traversal anyway, and this runs
;; on every syntax object the expander visits.
(define (syntax-search-everything stx)
  (define subforms-in-reverse '())
  (let loop ([stx stx])
    (set! subforms-in-reverse (cons stx subforms-in-reverse))
    (syntax-parse stx
      [(part ...)
       #:cut
       (for ([part-stx (in-list (attribute part))])
         (loop part-stx))]
      [(part ...+ . tail-part)
       #:cut
       (for ([part-stx (in-list (attribute part))])
         (loop part-stx))
       (loop #'tail-part)]
      [_ (void)]))
  (reverse subforms-in-reverse))


(module+ test
  (test-case "syntax-search"
    (define stx
      #'(define (foo)
          (cons x y)
          (define (bar)
            (cons a b))
          (cons c d)))
    (define actual
      (sequence->list
       (syntax-search stx
         #:literals (cons)
         [(cons _ _) (stream (syntax->datum this-syntax))])))
    (define expected '((cons x y) (cons a b) (cons c d)))
    (check-equal? actual expected)))


(define-syntax-parse-rule (syntax-find-first stx-expr
                            option:syntax-parse-option ...
                            syntax-pattern
                            directive:syntax-parse-pattern-directive ...)
  (let ()
    (define results (syntax-search stx-expr (~@ . option) ... [syntax-pattern (~@ . directive) ...]))
    (and (not (stream-empty? results)) (stream-first results))))


(module+ test
  (test-case "syntax-find-first"
    (define stx
      #'(define (foo)
          (cons x y)
          (define (bar)
            (cons a b))
          (cons c d)))
    (define actual (syntax-find-first stx #:literals (cons) #:datum-literals (a) (cons a _)))
    (check-equal? (syntax->datum actual) '(cons a b))
    (check-false (syntax-find-first stx #:literals (cons) (cons _ _ _)))))


(define (perform-syntax-traversal stx modifier
                                  #:skip-root? [skip-root? #false]
                                  #:parent-context-modifier [context-modifier* #false]
                                  #:parent-srcloc-modifier [srcloc-modifier (λ (_) #'here)]
                                  #:parent-props-modifier [props-modifier (λ (_) #false)])
  (define context-modifier
    (or context-modifier*
        (let ([scope (make-syntax-introducer)])
          (λ (stx) (scope stx 'add)))))
  (let loop ([stx stx] [root? #true])
    (define should-skip-because-root? (and skip-root? root?))
    (define modified-stx (and (not should-skip-because-root?) (modifier stx)))
    (match modified-stx
      [(== stx) stx]
      [(? syntax?) modified-stx]
      [#false
       (syntax-parse stx
         [(child ...)
          #:cut
          (define traversed-children
            (for/list ([child-stx (in-list (attribute child))])
              (loop child-stx #false)))
          (define all-same?
            (for/and ([traversed-child (in-list traversed-children)]
                      [child-stx (in-list (attribute child))])
              (equal? traversed-child child-stx)))
          (if all-same?
              stx
              (datum->syntax (context-modifier stx)
                             traversed-children
                             (srcloc-modifier stx)
                             (props-modifier stx)))]
         [(child ...+ . tail-child)
          #:cut
          (define traversed-children
            (for/list ([child-stx (in-list (attribute child))])
              (loop child-stx #false)))
          (define traversed-tail-child (loop (attribute tail-child) #false))
          (define all-same?
            (and (for/and ([traversed-child (in-list traversed-children)]
                           [child-stx (in-list (attribute child))])
                   (equal? traversed-child child-stx))
                 (equal? traversed-tail-child (attribute tail-child))))
          (if all-same?
              stx
              (datum->syntax (context-modifier stx)
                             (append traversed-children traversed-tail-child)
                             (srcloc-modifier stx)
                             (props-modifier stx)))]
         [_ stx])])))


(define-syntax-parse-rule
  (syntax-traverse (~var stx-expr (expr/c #'syntax?))
    (~optional (~seq #:skip-root? skip-root?))
    option:syntax-parse-option ...
    [clause-pattern directive:syntax-parse-pattern-directive ... clause-body:expr ...+] ...
    (~optional (~seq #:parent-context-modifier context-modifier:expr))
    (~optional (~seq #:parent-srcloc-modifier srcloc-modifier:expr))
    (~optional (~seq #:parent-props-modifier props-modifier:expr)))
    

  (let ()

    (define (traversal-case stx)
      (syntax-parse stx
        (~@ . option) ...
        [clause-pattern (~@ . directive) ... clause-body ...]
        ...
        [_ #false]))

    (perform-syntax-traversal stx-expr.c traversal-case
                              (~? (~@ #:skip-root? skip-root?))
                              (~? (~@ #:parent-context-modifier context-modifier))
                              (~? (~@ #:parent-srcloc-modifier srcloc-modifier))
                              (~? (~@ #:parent-props-modifier props-modifier)))))


(module+ test
  (test-case "syntax-traverse"
    (define stx
      #'(define (foo)
          (cons x y)
          (define (bar)
            (cons a b))
          (cons c d)))
    (define actual
      (syntax-traverse stx
        #:literals (cons)
        [(cons _ _) #'CONS-EXPRESSION]))
    (define expected
      '(define (foo)
         CONS-EXPRESSION
         (define (bar)
           CONS-EXPRESSION)
         CONS-EXPRESSION))
    (check-equal? (syntax->datum actual) expected))

  (test-case "syntax-traverse dotted tail syntax"
    (define stx #'(a b . c))
    (define actual (syntax-traverse stx [:id #'ID]))
    (define expected '(ID ID . ID))
    (check-equal? (syntax->datum actual) expected))

  (test-case "syntax-traverse #:skip-root? true"
    (define stx #'(a b (c d) e))
    (define actual
      (syntax-traverse stx
        #:skip-root? #true
        [(_ ...) #'LIST]))
    (check-equal? (syntax->datum actual) '(a b LIST e)))

  (test-case "syntax-traverse #:skip-root? true doesn't execute directives on root"
    (define stx #'(a b (c d) e))
    (define execution-count 0)
    (syntax-traverse stx
      #:skip-root? #true
      [(_ ...)
       #:do [(set! execution-count (add1 execution-count))]
       #'LIST])
    (check-equal? execution-count 1))

  (test-case "syntax-traverse #:skip-root? false"
    (define stx #'(a b (c d) e))
    (define actual
      (syntax-traverse stx
        #:skip-root? #false
        [(_ ...) #'LIST]))
    (check-equal? (syntax->datum actual) 'LIST))

  (test-case "syntax-traverse originality"
    (define stx (read-syntax #false (open-input-string "(1 2 (a b) 3 4)")))
    (check-true (syntax-original? stx))
    (define traversed-stx
      (syntax-traverse stx
        [(_ id:id) (attribute id)]))
    (check-equal? (syntax->datum traversed-stx) '(1 2 b 3 4))
    (check-false (syntax-original? traversed-stx))
    (check-not-equal? (syntax-srcloc traversed-stx) (syntax-srcloc stx))
    (define/syntax-parse (1* 2* b* 3* 4*) traversed-stx)
    (check-true (syntax-original? #'1*))
    (check-true (syntax-original? #'2*))
    (check-true (syntax-original? #'b*))
    (check-true (syntax-original? #'3*))
    (check-true (syntax-original? #'4*))))