brownplt/LambdaJS — Issue #4: Faster substitution function
Author: arjunguha
Created: 2012-01-02 12:08:20+00:00
State: open

```
(define-metafunction lambdaJS
subst-n : (x any) ... any -> any
[(subst-n (x any) ... any_body)
,(subst (term (x ...)) (term (any ...)) (term any_body))])

(define (subst vars replacements body)
(define replacements-ht
(for/fold ([m (hash)])
([v (in-list vars)]
[rep (in-list replacements)])
(hash-set m v rep)))
(define replacements-free-vars (for/list ([x (in-set (fvs replacements))]) x))
(define replacements-fresh-vars (variables-not-in (cons vars body)
replacements-free-vars))
(define init-fv-map
(for/fold ([m (hash)])
([fresh (in-list replacements-fresh-vars)]
[free (in-list replacements-free-vars)])
(hash-set m free fresh)))
(let loop ([body body]
[fvs-inactive init-fv-map]
[fvs-active (hash)]
[replacements replacements-ht])
(match body
[`(let ([,xs ,rhs] ...) ,body)
(define-values (new-xs new-inactive new-active new-replacements)
(adjust-active-inactive xs fvs-inactive fvs-active replacements))
`(let (,@(map (λ (x rhs) `[,x ,rhs]) new-xs rhs))
,(loop body new-inactive new-active new-replacements))]
[`(lambda (,xs ...) ,body)
(define-values (new-xs new-inactive new-active new-replacements)
(adjust-active-inactive xs fvs-inactive fvs-active replacements))
`(lambda (,@new-xs)
,(loop body new-inactive new-active new-replacements))]
[(? x? x)
(cond
[(hash-ref fvs-active x #f) => values]
[(hash-ref replacements x #f) => values]
[else x])]
[(? list?)
(map (λ (x) (loop x fvs-inactive fvs-active replacements))
body)]
[_ body])))

(define (adjust-active-inactive xs fvs-inactive fvs-active replacements)
(let loop ([xs xs]
[new-xs '()]
[fvs-inactive fvs-inactive]
[fvs-active fvs-active]
[replacements replacements])
(cond
[(null? xs)
(values (reverse new-xs)
fvs-inactive
fvs-active
replacements)]
[else
(define x (car xs))
(define inactive-var? (hash-has-key? fvs-inactive x))
(define active-var? (hash-has-key? fvs-active x))
(define new-x
(cond
[inactive-var? (hash-ref fvs-inactive x)]
[active-var? (hash-ref fvs-active x)]
[else x]))
(loop (cdr xs)
(cons new-x new-xs)
(if inactive-var?
(hash-remove fvs-inactive x)
fvs-inactive)
(if inactive-var?
(hash-set fvs-active x (hash-ref fvs-inactive x))
fvs-active)
(if (hash-has-key? replacements x)
(hash-remove replacements x)
replacements))])))

(define x? (redex-match lambdaJS x))

(define (fvs body)
(let loop ([body body])
(match body
[`(let ([,xs ,rhs] ...) ,body)
(set-subtract (loop body) (apply set xs))]
[`(lambda (,xs ...) ,body)
(set-subtract (loop body) (apply set xs))]
[(? x?)
(set body)]
[(? list?)
(for/fold ([fvs (set)])
([e (in-list body)])
(set-union fvs (loop e)))]
[_ (set)])))

(test-equal (fvs (term (+ x a b))) (set 'x 'a 'b))
(test-equal (fvs (term (lambda (x) (+ x y)))) (set 'y))
(test-equal (fvs (term (let ([x 1][y 2]) (+ x y z)))) (set 'z))

(test-equal (term (subst-n (x y) x)) (term y))
(test-equal (term (subst-n (x y) z)) (term z))
(test-equal (term (subst-n (x y) (x (y z)))) (term (y (y z))))
(test-equal (term (subst-n (x y) ((lambda (x) x) ((lambda (y1) y1)
(lambda (x) z)))))
(term ((lambda (x) x) ((lambda (y1) y1) (lambda (x) z)))))
(test-equal (term (subst-n (x y) (if0 (+ 1 x) x x)))
(term (if0 (+ 1 y) y y)))
(test-equal (term (subst-n (x (lambda (z) y)) (lambda (y) x)))
(term (lambda (y1) (lambda (z) y))))
(test-equal (term (subst-n (x 1) (lambda (y) x)))
(term (lambda (y) 1)))
(test-equal (term (subst-n (x y) (lambda (y) x)))
(term (lambda (y1) y)))
(test-equal (term (subst-n (x (lambda (y) y)) (lambda (z) (z2 z))))
(term (lambda (z) (z2 z))))
(test-equal (term (subst-n (x (lambda (z) z)) (lambda (z) (z1 z))))
(term (lambda (z) (z1 z))))
(test-equal (term (subst-n (x z) (lambda (z) (z1 z))))
(term (lambda (z2) (z1 z2))))
(test-equal (term (subst-n (x3 5) (lambda (x2) x2)))
(term (lambda (x2) x2)))
(test-equal (term (subst-n (z *) (lambda (z x) 1)))
(term (lambda (z x) 1)))
(test-equal (term (subst-n (q (lambda (x) z)) (lambda (z x) q)))
(term (lambda (z1 x) (lambda (x) z))))
(test-equal (term (subst-n (x 1) (lambda (x x) x)))
(term (lambda (x x) x)))
(test-equal (term (subst-n (x (y z)) (lambda (z) (z (x y)))))
(term (lambda (z1) (z1 ((y z) y)))))
(test-results)
```