[racket] Alpha equivalence checking for Racket

I've written a crude, incomplete and barely tested alpha equivalence
checker for Racket syntax in case anyone's interested. It should work
for most expression-level syntax. People have asked about
functionality like this before [1], so I figure I'd share my code. Is
this a good way to go about the problem or are there more concise
solutions?

[1] http://lists.racket-lang.org/dev/archive/2011-March/005958.html


#lang racket

(require syntax/parse
         syntax/id-table
         syntax/stx)

(define (alpha-equivalent? a b
                           ;; Descend into syntax-quote forms
                           #:do-syntax? [do-syntax? #t]
                           ;; Expand before checking alpha-equivalence
                           #:expand? [expand? #t])
  (norm-equal? (normalize a
                          #:do-syntax? do-syntax?
                          #:expand? expand?)
               (normalize b
                          #:do-syntax? do-syntax?
                          #:expand? expand?)))

(struct norm-id (i) #:transparent)
(define (normalize stx
                   #:do-syntax? [do-syntax? #t]
                   #:expand? [expand? #t])
  (define ctr 0)
  (define (inc) (set! ctr (add1 ctr)) ctr)
  (define env (make-free-id-table))
  (define (put new-id)
    (define new-nid (norm-id (inc)))
    (dict-set! env new-id new-nid)
    new-nid)
  (define (get id)
    (dict-ref env id id))
  (define (put-λ-formals formals)
    (syntax-parse #'formals
      [(arg ...+ . restid)
       `(,@(stx-map put #'(arg ...)) . ,(put #'restid))]
      [restid
       `(,(put #'restid))]))
  (define (norm stx)
    (syntax-parse stx
      [((~literal #%expression) expr)
       `(,#'#%expression ,(norm #'expr))]
      [((~literal begin-for-syntax) form ...)
       `(,#'begin-for-syntax ,(stx-map norm #'(form ...)))]
      [((~literal #%provide) raw-provide-spec ...) ;; TODO
       (error "Unimplemented" stx)]
      [((~literal #%declare) declaration-keyword ...)
       (error "Unimplemented" stx)] ;; TODO
      [((~literal module) id module-path
        ((~literal #%plain-module-begin) form ...))
       `(,#'module ,'id ,'module-path
          (,#'#%plain-module-begin ,@(stx-map norm #'(form ...))))]
      [((~literal module*) id module-path
        ((~literal #%plain-module-begin) form ...))
       `(,#'module* ,'id ,'module-path
          (,#'#%plain-module-begin ,@(stx-map norm #'(form ...))))]
      [((~and define (~or (~literal define-values)
                          (~literal define-syntaxes)))
        (id ...)
        expr)
       `(,#'define
          ,(stx-map put #'(id ...))
          ,(norm #'expr))]
      [((~literal #%require) raw-require-spec ...)
       (error "Unimplemented" stx)] ;; TODO
      [(~or d:keyword d:number d:boolean d:str d:char)
       (syntax->datum #'d)]
      [id:identifier
       (get #'id)]
      [((~literal #%plain-lambda) formals expr ...+)
       `(,#'#%plain-lambda
         ,@(put-λ-formals #'formals)
         ,@(stx-map norm #'(expr ...)))]
      [((~literal case-lambda) case ...)
       (for/list ([c (syntax->list #'(case ...))])
         (syntax-parse c
           [(formals expr ...+)
            `(,@(put-λ-formals #'formals)
              ,@(stx-map norm #'(expr ...)))]))]
      [((~literal if) p c a)
       `(,#'if ,@(stx-map norm #'(p c a)))]
      [((and begin~ (~or (~literal begin) (~literal begin0))) expr ...+)
       `(,#'begin~ ,@(stx-map norm #'(expr ...)))]
      [((~and let~ (~or (~literal let-values) (~literal letrec-values)))
        (bindings ...)
        expr
        ...+)
       (for/list ([b (syntax->list #'(bindings ...))])
         (syntax-parse b
           [[(id ...) _]
            (stx-map put #'(id ...))]))
       `(,#'let~
         ,(for/list ([b (syntax->list #'(bindings ...))])
            (syntax-parse b
              [[(id ...) expr]
               `(,(stx-map get #'(id ...)) ,(norm #'expr))]))
         ,@(stx-map norm #'(expr ...)))]
      [((~literal set!) id expr)
       `(,#'set! ,(get #'id) ,(norm #'expr))]
      [((~literal quote) datum)
       `(,#'quote ,(syntax->datum #'datum))] ;; FIXME?
      [(quote-syntax datum)
       `(,#'quote-syntax ,(if do-syntax? (norm #'datum) #'datum))]
      [((~literal with-continuation-mark) e1 e2 e3)
       `(,#'with-continuation-mark ,@(stx-map norm #'(e1 e2 e3)))]
      [((~literal #%plain-app) expr ...+)
       `(,#'#%plain-app ,@(stx-map norm #'(expr ...)))]
      [((~literal #%top) . id)
       `(,#'#%top . #'id)] ;; FIXME?
      [((~literal #%variable-reference) id)
       `(,#'#%variable-reference ,(get #'id))]
      [((~literal #%variable-reference) (#%top . id))
       `(,#'#%variable-reference (,#'#%top . #'id))] ;; FIXME?
      [((~literal #%variable-reference))
       `(,#'#%variable-reference)]
      [_ (error "Implementation error")]))
  (norm (if expand?
            (expand stx)
            stx)))

(define (norm-equal? a b)
  (match/values (values a b)
    [((list x ...) (list y ...))
     (if (= (length x) (length y))
         (andmap norm-equal? x y)
         #f)]
    [((? identifier?) (? identifier?))
     (free-identifier=? a b)]
    [(_ _) (equal? a b)]))