stxutils.rkt

#lang racket/base
(require
 (for-syntax
  racket/base
  syntax/parse
  racket/syntax)
 (for-template (only-in racket/base let-values #%plain-lambda))
 racket/function
 racket/match
 syntax/id-set
 syntax/parse
 racket/syntax
 syntax/stx
 syntax/parse/experimental/reflect)

(module+ test
  (require chk))

;; Tools for manipulating ASTs in Racket.
;; While In Theory (TM), we ought to be able to rely on hygiene and never worry
;; about binding when doing crazy macro magic, this is just not true.
;; So here are some tools for working with syntax objects when you can't rely on
;; hygiene and need to do it yourself.

;; A pattern for making these tools aware of binding:
;; open recursive binding parsers:

;; Racket plain-lambda parser
(define (plain-lambda-reconstructor xs _ es)
  #`(#%plain-lambda (#,@xs) #,@es))

;; ... if I only I had a maybe type
(define (plain-lambda-parser syn f)
  (syntax-parse syn
    #:literals (#%plain-lambda)
    [(#%plain-lambda (x) e)
     (list (list #'x) '() (list #'e) plain-lambda-reconstructor)]
    [_ (f syn f)]))

(module+ test
  (define matches1?
    (syntax-parser
      [binder
       #:when (plain-lambda-parser #'binder (lambda _ #f))
       #t]
      [_ #f]))
  (chk
   #:t (matches1? #`(#%plain-lambda (x) x))
   #:! #:t (matches1? #`(lambda (x) x))))

(define (lambda-reconstructor xs _ es)
  #`(lambda #,xs #,@es))

(define (lambda-parser syn f)
  (syntax-parse syn
    #:literals (lambda)
    [(lambda (x ...) e ...)
     (list (attribute x) '() (attribute e) lambda-reconstructor)]
    [_ (f syn f)]))

;; The current binding parser
(define current-binder-parser
  ;; Expects a syntax object and a recursive binder-parser (f), returns:
  ;; EITHER: #f, if it doesn't match
  ;; OR: a list containing:
  ;; - a list of ids
  ;; - a list of sub expressions in which ids are unbound
  ;; - a list of sub expressions in which ids are bound
  ;; - a reconstructor, which expects the (possibly modified) ids, unbounds
  ;;   exprs, and bound exprs and returns a syntax object representing the binding
  ;;   form of the same kind originally parsed.
  (make-parameter
   (lambda _ #f)
   ;; hide current-binder-parser's open recursive nature
   (lambda (f)
     (let ([old-f (current-binder-parser)])
       (case-lambda
         [(syn) (f syn old-f)]
         [(syn old-f) (f syn old-f)])))))

(module+ test
  (define matches2?
    (syntax-parser
      [binder
       #:when ((current-binder-parser) #'binder)
       #t]
      [_ #f]))

  (parameterize ([current-binder-parser plain-lambda-parser])
    (chk
     #:t (matches2? #`(#%plain-lambda (x) x))
     #:! #:t (matches2? #`(lambda (x) x)))
    (parameterize ([current-binder-parser lambda-parser])
      (chk
       #:t (matches2? #`(#%plain-lambda (x) x))
       #:t (matches2? #`(lambda (x) x))))))

;; Alpha equality

(define current-bound-vars (make-parameter (immutable-free-id-set)))

(define (id=? x y)
  (and (identifier? x) (identifier? y)
       (equal? (syntax->datum x) (syntax->datum y))))

;; Syntax equality, up to alpha equivalence.
;; Is aware of binding forms, and can be made to break hygiene by replacing id=?
;; with datum=?
(define (stx=? e1 e2)
  (cond
    [(id=? e1 e2)]
    [(and (number? (syntax-e e1)) (number? (syntax-e e2)))
     (= (syntax-e e1) (syntax-e e2))]
    [(and (stx-pair? e1) (stx-pair? e2))
     (syntax-parse (list e1 e2)
       [(binder1 binder2)
        #:when ((current-binder-parser) #'binder1)
        #:when ((current-binder-parser) #'binder2)
        (match (cons ((current-binder-parser) #'binder1)
                     ((current-binder-parser) #'binder2))
          [(cons (list xs1 ues1 bes1 r1)
                 (list xs2 ues2 bes2 r2))
           (and 
            ;; short-circuit
            (= (length xs1) (length xs2))
            (= (length ues1) (length ues2))
            (= (length bes1) (length bes2))

            (andmap stx=? ues1 ues2)
            (andmap stx=? bes1 (map (lambda (syn) (subst* xs1 xs2 syn)) bes2)))])]
       [_
        (and
         (= (length (syntax->list e1)) (length (syntax->list e2)))
         (andmap stx=? (syntax->list e1) (syntax->list e2)))])]
    [else #f]))

(module+ test
  (chk
   #:t (stx=? #'5 #'5)
   #:! #:t (stx=? #'5 #'4)
   #:! #:t (stx=? #'(lambda (x) x) #'(lambda (y) x))
   #:! #:t (stx=? #'(lambda (x) x) #'(lambda (y) y)))
  (parameterize ([current-binder-parser lambda-parser])
    (chk #:t (stx=? #'(lambda (x) x) #'(lambda (y) y)))))

;; Generic AST substitution
;; Replace the nodes stx=? to e2 by e1 in syn, unless e2 is a bound variable in bvs
;; Relies on current-binder-parser to detect binding forms
(define (subst e1 e2 syn)
  (let subst ([syn syn])
    (syntax-parse syn
      [e
       #:when (and (stx=? #'e e2)
                   (or (not (identifier? #'e))
                       (not (free-id-set-member? (current-bound-vars) #'e))))
       e1]
      ; something wrong with hygiene, so need this explicit case
      ; o.w., λ binders get substed
      [binder
       #:when ((current-binder-parser) #'binder)
       (match ((current-binder-parser) #'binder)
         [(list xs ues bes r)
          (r xs
             (map subst ues)
             (parameterize ([current-bound-vars (for/fold ([bvs (current-bound-vars)])
                                                          ([x xs])
                                                  (free-id-set-add bvs x))])
               (map subst bes)))])]
      [(e ...)
       (datum->syntax syn (map subst (attribute e)))]
      [_ syn])))

;; find for ASTs
;; returns the node stx=? to e0 in syn, if it exists.
;; otherwise, returns #f
(define (find-in e0 stx)
  (syntax-parse stx
    [e #:when (stx=? #'e e0) #'e]
    [(e ...)
     (for/first ([e (syntax->list #'(e ...))]
                 #:when (find-in e0 e))
       (find-in e0 e))]
    [_ #f]))

;; takes a list of values and a list of identifiers, in dependency order, and substitutes them into syn.
;; TODO PERF: reverse
(define (subst* v-ls x-ls syn)
  (for/fold ([syn syn])
            ([v (reverse v-ls)]
             [x (reverse x-ls)])
    (subst v x syn)))