#lang plai/gc2/collector
(define (init-allocator)
  (heap-set! 0 1))

(define (malloc n)
  (define addr (heap-ref 0))
  (unless (<= (+ addr n) (heap-size))
    (error 'allocator "out of memory"))
  (heap-set! 0 (+ addr n))
  addr)

(define (gc:flat? addr)
  (equal? (heap-ref addr) 'flat))

(define (gc:alloc-flat x)
  (define addr (malloc 2))
  (heap-set! addr 'flat)
  (heap-set! (+ addr 1) x)
  addr)

(define (gc:deref addr)
  (unless (equal? (heap-ref addr) 'flat)
    (error 'gc:deref "not a flat at ~a" addr))
  (heap-ref (+ addr 1)))

(define (gc:cons f r)
  (define addr (malloc 3))
  (heap-set! addr 'cons)
  (heap-set! (+ addr 1) (read-root f))
  (heap-set! (+ addr 2) (read-root r))
  addr)

(define (gc:cons? addr)
  (equal? (heap-ref addr) 'cons))

(define (gc:first addr)
  (check-pair addr)
  (heap-ref (+ addr 1)))

(define (gc:rest addr)
  (check-pair addr)
  (heap-ref (+ addr 2)))

(define (check-pair addr)
  (unless (equal? (heap-ref addr) 'cons)
    (error 'check-pair "not a pair @ ~a" addr)))

(define (gc:set-first! addr v)
  (check-pair addr)
  (heap-set! (+ addr 1) v))
(define (gc:set-rest! addr v)
  (check-pair addr)
  (heap-set! (+ addr 2) v))

(define (gc:closure code-pointer free-vars)
  (define addr (malloc (+ 2 (length free-vars))))
  (heap-set! addr 'clos)
  (heap-set! (+ addr 1) code-pointer)
  (for ([i (in-range 0 (length free-vars))]
        [v (in-list free-vars)])
    (heap-set! (+ addr 2 i)
               (read-root v)))
  addr)

(define (gc:closure? addr)
  (equal? (heap-ref addr) 'clos))
(define (gc:closure-code-ptr addr)
  (unless (gc:closure? addr)
    (error "not a closure @ ~a" addr))
  (heap-ref (+ addr 1)))
(define (gc:closure-env-ref addr i)
  (unless (gc:closure? addr)
    (error "not a closure @ ~a" addr))
  (heap-ref (+ addr 2 i)))

(module+ test
  (with-heap (vector 'x 'x 'x 'x 'x)
    (init-allocator)
    (gc:alloc-flat #f)
    (test (current-heap) (vector 3 'flat #f 'x 'x))))

(module+ test
  (with-heap (vector 'x 'x 'x 'x 'x 'x 'x 'x 'x)
    (init-allocator)
    (gc:cons
     (simple-root (gc:alloc-flat #f))
     (simple-root (gc:alloc-flat #t)))
    (test (current-heap)
          (vector 8 'flat #f 'flat #t 'cons 1 3 'x))))