census.rkt

#lang racket

(define variable? symbol?)
(define (abstraction? t) (and (pair? t) (eq? (car t) 'lambda)))
(define (branch? t) (and (pair? t) (eq? (car t) 'if)))
(define (call/cc? t) (and (pair? t) (eq? (car t) 'call/cc)))

(struct acont (f xs vs r c))
(struct fcont (vs c))
(struct bcont (conseq altern r c))
(struct ccont (c))

(struct closure (t r))
(struct cont (k))

(define (map* f xs)
  (cond [(null? xs) '()]
        [(pair? xs) (cons (f (car xs)) (map* f (cdr xs)))]
        [else (f xs)]))

(define (extend r xs vs)
  (cond [(null? xs)
         (cond [(null? vs) r]
               [else (error "mismatch:" xs vs)])]
        [(symbol? xs) (cons (cons xs vs) r)]
        [else (cons (cons (car xs) (car vs))
                    (extend r (cdr xs) (cdr vs)))]))
(define (lookup r x)
  (cond [(assq x r) => cdr]
        [else (error "not found:" x)]))

(define (cps t)
  (let ([k (gensym 'halt)])
    `(lambda (,k)
       ,(cps/t t '() k))))

(define (cps/t t r c)
  (cond [(variable? t) (ret c (lookup r t))]
        [(abstraction? t) (ret c (closure t r))]
        [(branch? t) (cps/t (cadr t) r (bcont (caddr t) (cadddr t) r c))]
        [(call/cc? t) (cps/t (cadr t) r (ccont c))]
        [else (cps/call (car t) (cdr t) '() r c)]))

(define (cps/call f xs vs r c)
  (if (null? xs)
      (cps/t f r (fcont vs c))
      (cps/t (car xs) r (acont f (cdr xs) vs r c))))

(define (ret c v) (ret* c (list v)))

(define (ret* c vs)
  (cond [(symbol? c)
         (let loop ([as '()] [vs vs])
           (cond [(null? vs)
                  `(,(C c) ,@(reverse as))]
                 [(pair? vs)
                  (loop (cons (V (car vs)) as) (cdr vs))]
                 [else
                  `(apply ,(C c) ,@(reverse (cons (V vs) as)))]))]
        [(acont? c)
         (if (= (length vs) 1)
             (cps/call (acont-f c) (acont-xs c) (cons (car vs) (acont-vs c))
                       (acont-r c) (acont-c c))
             (error "acont expects exact 1 val:" vs))]
        [(fcont? c)
         (if (= (length vs) 1)
             (call (car vs) (fcont-c c) (reverse (fcont-vs c)))
             (error "fcont expects exact 1 val:" vs))]
        [(bcont? c)
         (if (= (length vs) 1)
             (let ([k (gensym 'k)])
               `((lambda (,k)
                   (if ,(car vs)
                       ,(cps/t (bcont-conseq c) (bcont-r c) k)
                       ,(cps/t (bcont-altern c) (bcont-r c) k)))
                 ,(C (bcont-c c))))
             (error "bcont expects exact 1 val:" vs))]
        [(ccont? c)
         (if (= (length vs) 1)
             (cond [(symbol? (ccont-c c))
                    (let ([k (ccont-c c)])
                      (call (car vs) k (list (cont k))))]
                   [else
                    (let ([k (gensym 'k)])
                      `((lambda (,k)
                          ,(call (car vs) k (list (cont k))))
                        ,(C (ccont-c c))))])
             (error "ccont execpts exact 1 val:" vs))]
        [else (error "unknown cont:" c)]))

(define (call f c vs)
  (cond [(closure? f)
         (let ([args (cadr (closure-t f))]
               [body (caddr (closure-t f))]
               [r (closure-r f)])
           (inline body r args vs c))]
        [(symbol? f) `(,(V f) ,(C c) ,@(map V vs))]
        [else (error "unknown operator" f)]))

(define (inline body r xs vs c)
  (let f ([xs xs] [vs vs] [r r])
    (cond [(null? vs)
           (if (null? xs)
               (cps/t body r c)
               (error "not match:" xs))]
          [(symbol? (car vs))
           (f (cdr xs) (cdr vs)
              (extend r (list (car xs)) (list (car vs))))]
          [(or (closure? (car vs)) (cont? (car vs)))
           (let ([y (gensym 'y)])
             `((lambda (,y)
                 ,(f (cdr xs) (cdr vs)
                     (extend r (list (car xs)) (list y))))
               ,(V (car vs))))]
          [else (error "unknown value" (car vs))])))

(define (V v)
  (cond [(symbol? v) v]
        [(closure? v)
         (let* ([k (gensym 'k)]
                [xs (cadr (closure-t v))]
                [ys (map* (lambda (x) (gensym x)) xs)])
           `(lambda (,k ,@ys)
              ,(cps/t (caddr (closure-t v)) (extend (closure-r v) xs ys) k)))]
        [(cont? v)
         (let ([k (gensym 'k)] [xs (gensym 'xs)])
           `(lambda (,k . ,xs)
              ,(ret* (cont-k v) xs)))]
        [else (error "unknown value" v)]))

(define (C c)
  (cond [(symbol? c) c]
        [(or (acont? c)
             (fcont? c)
             (bcont? c)
             (ccont? c))
         (let ([x (gensym 'x)])
           `(lambda (,x) ,(ret c x)))]
        [else (error "unknown cont:" c)]))




(define h (make-hasheq))
(define (extend! x b) (hash-set! h x b))
(define (update! x b) (hash-set! h x b))
(define (lookup2 x)
  (let ([b (hash-ref h x)])
    (cond [(and (top? b) (variable? (top-v b)))
           (lookup2 (top-v b))]
          [(and (single? b) (variable? (single-v b)))
           (lookup2 (single-v b))]
          [else b])))

(struct top (v))
(struct single (v))
(struct bottom ())

(define (census/v v)
  (cond [(variable? v)
         (let ([b (lookup2 v)])
           (update! v (bottom))
           (cond [(top? b) (census/v (top-v b))]
                 [(single? b)
                  (cond [(variable? (single-v b)) (census/v (single-v b))]
                        [else (for-each census/v (cadr (single-v b)))])]))]
        [(abstraction? v)
         (let ([xs (cadr v)] [m (caddr v)])
           (for-each (lambda (x) (extend! x (bottom))) xs)
           (census/m (car m) (cdr m)))]))

(define (census/m f as)
  (cond [(variable? f)
         (let ([b (lookup2 f)])
           (cond [(top? b)
                  (update! f (single (top-v b)))
                  (for-each (lambda (x a) (extend! x (top a)))
                            (cadr (top-v b))
                            as)
                  (let ([m (caddr (top-v b))])
                    (census/m (car m) (cdr m)))]
                 [(single? b)
                  (update! f (bottom))
                  (for-each census/v (cadr (single-v b)))
                  (for-each (lambda (v) (census/v v)) as)]
                 [(bottom? b)
                  (for-each (lambda (v) (census/v v)) as)]))]
        [(abstraction? f)
         (let ([xs (cadr f)] [m (caddr f)])
           (for-each (lambda (x v) (extend! x (top v))) xs as)
           (census/m (car m) (cdr m)))]))

simplify.ss

#!r6rs

(import (rnrs) (rnrs mutable-pairs) (srfi :78) (felis match))

(define gensym
  (let ([n 0])
    (lambda (name)
      (set! n (+ n 1))
      (string->symbol
       (string-append name "@" (number->string n))))))

(define (cps/k e k)
  (match e
    [(? symbol? x) `(,k ,x)]
    [()
     (let ([x (gensym "x")])
       `(letval ,x = () in (,k ,x)))]
    [`(cons ,e1 ,e2)
     (let ([p (gensym "p")])
       (cps/c e1 (lambda (z1)
                   (cps/c e2 (lambda (z2)
                               `(letval ,p = (cons ,z1 ,z2) in (,k ,p)))))))]
    [`(in ,i ,e)
     (let ([x (gensym "x")])
       (cps/c e (lambda (z)
                  `(letval ,x = (in ,i ,z) in (,k ,x)))))]
    [`(pi ,i ,e)
     (let ([x (gensym "x")])
       (cps/c e (lambda (z)
                  `(let ,x = (pi ,i ,z) in (,k ,x)))))]
    [`(case ,e of [(in 1 ,x1) => ,e1] [(in 2 ,x2) => ,e2])
     (let ([k1 (gensym "k1")]
           [k2 (gensym "k2")])
       (cps/c e (lambda (z)
                  `(letcont ,k1 ,x1 = ,(cps/k e1 k) in
                            (letcont ,k2 ,x2 = ,(cps/k e2 k) in
                                     (case ,z of ,k1 ,k2))))))]
    [`(let val ,x = ,e1 in ,e2)
     (let ([j (gensym "j")])
       `(letcont ,j ,x = ,(cps/k e2 k) in ,(cps/k e1 j)))]
    [`(fn ,x => ,e)
     (let ([f (gensym "f")]
           [j (gensym "j")])
       `(letval ,f = (lambda (,j ,x) ,(cps/k e j)) in (,k ,f)))]
    [`(,e1 ,e2)
     (cps/c e1 (lambda (z1)
                 (cps/c e2 (lambda (z2)
                             `(,z1 ,k ,z2)))))]))

(define (cps/c e c)
  (match e
    [(? symbol? x) (c x)]
    [()
     (let ([x (gensym "x")])
       `(letval ,x = () in ,(c x)))]
    [`(cons ,e1 ,e2)
     (let ([p (gensym "p")])
       (cps/c e1 (lambda (z1)
                   (cps/c e2 (lambda (z2)
                               `(letval ,p = (cons ,z1 ,z2) in ,(c p)))))))]
    [`(in ,i ,e)
     (let ([x (gensym "x")])
       (cps/c e (lambda (z)
                  `(letval ,x = (in ,i ,z) in ,(c x)))))]
    [`(pi ,i ,e)
     (let ([x (gensym "x")])
       (cps/c e (lambda (z)
                  `(let ,x = (pi ,i ,z) in ,(c x)))))]
    [`(case ,e of [(in 1 ,x1) => ,e1] [(in 2 ,x2) => ,e2])
     (let ([k (gensym "k")]
           [x (gensym "x")]
           [k1 (gensym "k1")]
           [k2 (gensym "k2")])
       (cps/c e (lambda (z)
                  `(letcont ,k ,x = ,(c x) in
                            (letcont ,k1 ,x1 = ,(cps/k e1 k) in
                                     (letcont ,k2 ,x2 = ,(cps/k e2 k) in
                                              (case ,z of ,k1 ,k2)))))))]
    [`(fn ,x => ,e)
     (let ([f (gensym "f")]
           [j (gensym "j")])
       `(letval ,f = (lambda (,j ,x) ,(cps/k e j)) in ,(c f)))]
    [`(,e1 ,e2)
     (let ([k (gensym "k")]
           [x (gensym "x")])
       (cps/c e1 (lambda (z1)
                   (cps/c e2 (lambda (z2)
                               `(letcont ,k ,x = ,(c x) in (,z1 ,k ,z2)))))))]))

(define (probe census K)
  (match K
    [`(letcont ,k ,x = ,L in ,K)
     (probe (probe census L) K)]
    [`(letval ,x = ,V in ,K)
     (probe (probeV census V) K)]
    [`(let ,x = (pi ,i ,y) in ,K)
     (probe (add census y) K)]
    [`(case ,x of ,k1 ,k2) (add (add (add census x) k1) k2)]
    [`(,k ,x) (add (add census k) x)]
    [`(,f ,k ,x) (add (add (add census f) k) x)]))

(define (probeV census V)
  (match V
    [() census]
    [`(cons ,x ,y) (add (add census x) y)]
    [`(in ,i ,x) (add census x)]
    [`(lambda (,k ,x) ,K) (probe census K)]))

(define (add census x) (cons x census))

(define (count census x)
  (length (filter (lambda (y) (eq? x y)) census)))

(define (simplify census env S K)
  (match K
    [`(letcont ,k ,x = ,L in ,K)
     (if (= (count census k) 0)
         (simplify census env S K)
         (match L
           [`(,j ,y)
            (=> fail)
            (if (eq? x y)
                (simplify census env (addS S k j) K)
                (fail))]
           [_
            `(letcont ,k ,x = ,(simplify census env S L) in
                      ,(simplify census (extend env k `(lambda (,x) ,L)) S K))]))]
    [`(letval ,x = ,V in ,K)
     (if (= (count census x) 0)
         (simplify census env S K)
         (match V
           [`(lambda (,k ,y) (,f ,j ,z))
            (=> fail)
            (if (and (eq? k j) (eq? y z))
                (simplify census env (addS S x f) K)
                (fail))]
           [`(cons ,z1 ,z2)
            (=> fail)
            (match (list (lookup env z1) (lookup env z2))
              [`((pi 1 ,p1) (pi 2 ,p2))
               (=> fail)
               (if (eq? p1 p2)
                   (simplify census env (addS S x p1) K)
                   (fail))]
              [_ (fail)])]
           [_
            `(letval ,x = ,(simplifyV census env S V) in
                     ,(simplify census (extend env x V) S K))]))]
    [`(let ,x = (pi ,i ,y) in ,K)
     (let ([y (applyS S y)])
       (match (lookup env y)
         [`(cons ,z1 ,z2)
          (simplify census env (addS S x (case i [(1) z1] [(2) z2])) K)]
         [_
          `(let ,x = (pi ,i ,y) in
             ,(simplify census (extend env x `(pi ,i ,y)) S K))]))]
    [`(case ,x of ,k1 ,k2)
     (let ([x (applyS S x)] [k1 (applyS S k1)] [k2 (applyS S k2)])
       (match (lookup env x)
         [`(in ,i ,y) `(,(case i [(1) k1] [(2) k2]) ,y)]
         [_
          (match (list (lookup env k1) (lookup env k2))
            [`((lambda (,x1) (letval ,y1 = (in 1 ,xx1) in (,j1 ,yy1)))
               (lambda (,x2) (letval ,y2 = (in 2 ,xx2) in (,j2 ,yy2))))
             (=> fail)
             (if (and (eq? x1 xx1) (eq? y1 yy1)
                      (eq? x2 xx2) (eq? y2 yy2)
                      (eq? j1 j2))
                 `(,j1 ,x)
                 (fail))]
            [_ `(case ,x of ,k1 ,k2)])]))]
    [`(,f ,k ,x)
     (let ([f (applyS S f)] [k (applyS S k)] [x (applyS S x)])
       (if (= (count census f) 1)
           (match (lookup env f)
             [`(lambda (,j ,y) ,K)
              (simplify census env (addS (addS S j k) y x) K)]
             [_
              `(,f ,k ,x)])
           `(,f ,k ,x)))]
    [`(,k ,x)
     (let ([k (applyS S k)] [x (applyS S x)])
       (if (= (count census k) 1)
           (match (lookup env k)
             [`(lambda (,y) ,K)
              (simplify census env (addS S y x) K)]
             [_
              `(,k ,x)])
           `(,k ,x)))]))

(define (simplifyV census env S V)
  (match V
    [() '()]
    [`(cons ,x ,y) `(cons ,(applyS S x) ,(applyS S y))]
    [`(in ,i ,x)   `(in ,i ,(applyS S x))]
    [`(lambda (,k ,x) ,K)
     `(lambda (,k ,x) ,(simplify census env S K))]
    [`(lambda (,x) ,K)
     `(lambda (,x) ,(simplify census env S K))]))

(define (extend env x V)
  (cons (cons x V) env))

(define (lookup env x)
  (let ([t (assq x env)])
    (and t (cdr t))))

(define (addS S x y)
  (cons (cons x y) S))

(define (applyS S x)
  (let ([t (assq x S)])
    (if t
        (cdr t)
        x)))

;; tests

(define e `(pi 1 ((fn x => (cons (g x) x)) y)))
(define K0 (cps/k e 'halt))
(define census0 (probe '() K0))

(define K1 (simplify census0 '() '() K0))
(define census1 (probe '() K1))

(define K2 (simplify census1 '() '() K1))
(define census2 (probe '() K2))

(define K3 (simplify census2 '() '() K2))
(define census3 (probe '() K3))

(define K4 (simplify census3 '() '() K3))

(check K4 => '(g halt y))

(check-report)