#lang plai/gc2/collector
;; Originally by Yixi Zhang
;; simplified for class by db

(require racket/stream)

;; config for collection
(define step-length 10)
(define (1st-gen-size) (round (* (heap-size) 1/4)));  size of young generation
(define (2nd-gen-end) (heap-size)) ; size of old generation
(define (2nd-gen-start)
  (define (table-size) (round (* (heap-size) 1/8)))
  (+ (1st-gen-size) (table-size)));  start of old generation
(define (2nd-gen-alloc-start)  (+ 4 (2nd-gen-start))) ; start position to allocate objects in old generation
(define (alloc-word) 0) ; start of objects allocation for young generation
(define (status-word) (2nd-gen-start)) ; current status of old generation incremental gc: in/out"
(define (free-list-head)   (+ 1 (2nd-gen-start))) ; head of free list
(define (step-count-word) (+ 2 (2nd-gen-start))) ; work count for each tracing round
(define (tracing-head-word) (+ 3 (2nd-gen-start))) ;head of tracing tree
(define (table-start-word) (1st-gen-size)) ; "start position of table for cross-generation pointers"

(define debug #t)
(define (print-heap-layout)
  (when debug
    (printf "init allocator 1-gen ~a 2-gen ~a ~a ~a fl ~a\n" (1st-gen-size)
            (2nd-gen-start) (2nd-gen-end) (2nd-gen-alloc-start) (free-list-head))))

(define (init-allocator)
  (unless (= 0 (modulo (heap-size) 256))
    (error 'init-allocator "heap size is not multiple of 256"))
  (for ([i (in-range 0 (heap-size))])
    (heap-set! i 'free)); ⋮
  (heap-set! (alloc-word) 1)
  (heap-set! (status-word) 'out)
  (let ([s (2nd-gen-alloc-start)])
    (heap-set! s 'free-n)
    (heap-set! (+ 1 s) #f)
    (heap-set! (+ 2 s) (- (2nd-gen-end) s))
    (heap-set! (free-list-head) s))
  (heap-set! (step-count-word) 0)
  (heap-set! (tracing-head-word) #f)
  (heap-set! (table-start-word)
             (+ (table-start-word) 1))
  (print-heap-layout))

(module+ test
  (with-heap (make-vector 512)
    (init-allocator)))

;; gc:deref : loc -> heap-value
;; must signal an error if fl-loc doesn't point to a flat value
(define (gc:deref fl-loc)
  (cond
    [(gc:flat? fl-loc) (heap-ref (+ fl-loc 1))]
    [else (error 'gc:deref "non-flat @ ~s" fl-loc)]))

;; track/loc : loc -> loc
;; if loc points to a flat or pair or proc, then return loc
;; else if loc points to a frwd, return the frwd address
(define (track/loc loc)
  (case (heap-ref loc)
    [(free) (error 'track/loc "wrong tag ~s @ ~a" (heap-ref loc) loc)]
    [(frwd) (heap-ref (+ loc 1))]
    [else loc]))

;; gc:alloc-flat : heap-value -> loc
(define (gc:alloc-flat fv)
  (define ptr (malloc 2))
  (heap-set! ptr 'flat)
  (heap-set! (+ ptr 1) fv)
  ptr)

;; ->location : (or/c location? root?) . -> . location?
(define (->location thing)
  (cond
    [(location? thing) thing]
    [(root? thing) (read-root thing)]))

;; gc:cons : loc loc -> loc
;; hd and tl are guaranteed to have been earlier
;; results from either gc:alloc-flat or gc:cons
(define (gc:cons hd tl)
  (define ptr (malloc 3 hd tl))
  (define hd/loc (->location hd))
  (define tl/loc (->location tl))
  (define head (track/loc hd/loc))
  (define tail (track/loc tl/loc))
  (when (and (= ptr 1)
             (or (need-forwarding-pointers? hd/loc)
                 (need-forwarding-pointers? tl/loc)))
    (free-1st-gen))
  (heap-set! ptr 'pair)
  (heap-set! (+ ptr 1) head)
  (heap-set! (+ ptr 2) tail)
  ptr)

;; gc:first : loc -> loc
;; must signal an error of pr-loc does not point to a pair
(define (gc:first pr-loc)
  (if (gc:cons? pr-loc)
      (heap-ref (+ (track/loc pr-loc) 1))
      (error 'first "non pair @ ~s" pr-loc)))

;; gc:rest : loc -> loc
;; must signal an error of pr-loc does not point to a pair
(define (gc:rest pr-loc)
  (if (gc:cons? pr-loc)
      (heap-ref (+ (track/loc pr-loc) 2))
      (error 'rest "non pair @ ~s" pr-loc)))

;; gc:flat? : loc -> boolean
;; loc is guaranteed to have been an earlier
;; result from either gc:alloc-flat or gc:cons
(define (gc:flat? loc)
  (case (heap-ref loc)
    [(flat grey-flat white-flat) #t]
    [(frwd) (gc:flat? (heap-ref (+ loc 1)))]
    [else #f]))

;; gc:cons? : loc -> boolean
;; loc is guaranteed to have been an earlier
;; result from either gc:alloc-flat or gc:cons
(define (gc:cons? loc)
  (case (heap-ref loc)
    [(pair grey-pair white-pair) #t]
    [(frwd) (gc:cons? (heap-ref (+ loc 1)))]
    [else #f]))

;; gc:set-first! : loc loc -> void
;; must signal an error of pr-loc does not point to a pair
(define (gc:set-first! pr-loc new)
  (cond
    [(gc:cons? pr-loc)
     (define loc (track/loc pr-loc))
     (heap-set! (+ loc 1) new)
     (when (and (2nd-gen? loc)
                (1st-gen? new))
       (table/alloc (+ loc 1) new))]
    [else (error 'set-first! "non pair @ ~s" pr-loc)]))

;; gc:set-rest! : loc loc -> void
;; must signal an error of pr-loc does not point to a pair
(define (gc:set-rest! pr-loc new)
  (cond
    [(gc:cons? pr-loc)
     (define loc (track/loc pr-loc))
     (heap-set! (+ loc 2) new)
     (when (and (2nd-gen? loc)
                (1st-gen? new))
       (table/alloc (+ loc 2) new))]
    [else (error 'set-rest! "non pair @ ~s" pr-loc)]))

;; gc:closure : heap-value (listof loc) -> loc
;; allocates a closure with 'code-ptr' and the free variables
;; in the vector 'free-vars'.
(define (gc:closure code-ptr free-vars)
  (define fv-count (length free-vars))
  (define next (malloc (+ fv-count 3)
                       free-vars
                       '()))
  (define updated-free-vars (for/vector ([v (in-list free-vars)])
                              (track/loc (read-root v))))
  (when (and (= next 1)
             (need-forwarding-pointers? free-vars))
    (free-1st-gen))
  (heap-set! next 'proc)
  (heap-set! (+ next 1) code-ptr)
  (heap-set! (+ next 2) fv-count)
  (for ([x (in-range 0 fv-count)])
    (heap-set! (+ next 3 x)
               (vector-ref updated-free-vars x)))
  next)

;; gc:closure-code-ptr : loc -> heap-value
;; given a location returned from an earlier allocation
;; check to see if it is a closure; if not signal an
;; error. if so, return the code-ptr
(define (gc:closure-code-ptr loc)
  (if (gc:closure? loc)
      (heap-ref (+ (track/loc loc) 1))
      (error 'gc:closure-code-ptr "non closure at ~a" loc)))

;; gc:closure-env-ref : loc number -> loc
;; given a location returned from an earlier allocation
;; check to see if it is a closure; if not signal an
;; error. if so, return the 'i'th variable in the closure
(define (gc:closure-env-ref loc i)
  (if (gc:closure? loc)
      (heap-ref (+ (track/loc loc) 3 i))
      (error 'gc:closure-env-ref "non closure at ~a" loc)))

;; gc:closure? : loc -> boolean
;; determine if a previously allocated location
;; holds a closure
(define (gc:closure? loc)
  (case (heap-ref loc)
    [(proc grey-proc white-proc) #t]
    [(frwd) (gc:closure? (heap-ref (+ loc 1)))]
    [else #f]))

(define (table/alloc pointer target)
  (define next (heap-ref (table-start-word)))
  (cond
    [(>= (+ next 2) (heap-size))
     (move/pointers (+ 1 (table-start-word)))
     (heap-set! (+ 1 (table-start-word)) pointer)
     (heap-set! (+ 2 (table-start-word)) target)
     (heap-set! (table-start-word) (+ 3 (table-start-word)))]
    [else
     (heap-set! next pointer)
     (heap-set! (+ next 1) target)
     (heap-set! (table-start-word) (+ next 2))]))

;; malloc : number[size] roots -> loc
(define (malloc n . extra-roots)
  (define addr (heap-ref (alloc-word)))
  (cond
    [(space-on-young-heap? addr n)
     (heap-set! (alloc-word) (+ addr n))
     addr]
    [else
     (collect-garbage extra-roots)
     (unless (need-forwarding-pointers? extra-roots)
       (free-1st-gen))
     (unless (space-on-young-heap? 1 n)
       (error 'alloc "object too large"))
     (heap-set! (alloc-word) (+ 1 n))
     1]))

(define (space-on-young-heap? start size)
  (<= (+ start size) (1st-gen-size)))

(define (need-forwarding-pointers? thing)
  (cond
    [(list? thing) (ormap need-forwarding-pointers? thing)]
    [(root? thing) (1st-gen? (read-root thing))]
    [(number? thing) (1st-gen? thing)]
    [else thing]))

;; collect-garbage : roots roots -> void
(define (collect-garbage . extra-roots)
  ;; young->old live objects copying
  (heap-set! (status-word) 'in)
  (make-pointers-to-2nd-gen-roots 1)
  (traverse/roots (get-root-set))
  (traverse/roots extra-roots)
  (move/pointers (+ 1 (table-start-word)))
  ;; only free/mark-white! if traversal is done
  (unless (heap-ref (tracing-head-word))
    (free/mark-white!
     (2nd-gen-alloc-start) #f #f #f))
  (heap-set! (status-word) 'out))

(define (traverse/roots thing)
  (cond
    [(list? thing)
     (for-each traverse/roots thing)]
    [(root? thing)
     (cond
       [(1st-gen? (read-root thing)) (move/roots thing)]
       [else (trace/roots-iff-white thing)])]
    [(number? thing)
     (cond
       [(1st-gen? thing) (move/ref (move/loc thing))]
       [else (when (white? thing) (push/cont thing))])]))

;; move/roots : loc/(listof loc) -> loc
;; move every thing reachable from 'roots'
;; to the to space
(define (move/roots thing)
  (cond
    [(list? thing)
     (for-each move/roots thing)]
    [(root? thing)
     (define new-addr (move/loc (read-root thing)))
     (set-root! thing new-addr)
     (move/ref new-addr)]
    [(number? thing)
     (move/ref (move/loc thing))]))

;; move/loc : loc -> loc
;; move object to the other semi-space
;; and return the new addr of moved object
(define (move/loc loc)
  (cond
    [(1st-gen? loc)
     (case (heap-ref loc)
       [(flat)
        (define new-addr (copy/alloc 2 #f #f))
        (heap-set! new-addr 'flat)
        (heap-set! (+ new-addr 1) (heap-ref (+ loc 1)))
        (heap-set! loc 'frwd)
        (heap-set! (+ loc 1) new-addr)
        new-addr]
       [(pair)
        (define new-addr (copy/alloc 3
                                     (heap-ref (+ loc 1))
                                     (heap-ref (+ loc 2))))
        (heap-set! new-addr 'pair)
        (heap-set! (+ new-addr 1) (track/loc (heap-ref (+ loc 1))))
        (heap-set! (+ new-addr 2) (track/loc (heap-ref (+ loc 2))))
        (heap-set! loc 'frwd)
        (heap-set! (+ loc 1) new-addr)
        new-addr]
       [(proc)
        (define length (+ 3 (heap-ref (+ loc 2))))
        (define free-vars (build-vector (- length 3)
                                        (lambda (i)
                                          (heap-ref (+ loc 3 i)))))
        (define new-addr (copy/alloc length free-vars '()))
        (heap-set! new-addr 'proc)
        (heap-set! (+ 1 new-addr) (heap-ref (+ 1 loc)))
        (heap-set! (+ 2 new-addr) (heap-ref (+ 2 loc)))
        (for ([x (in-range 3 length)])
          (heap-set! (+ new-addr x) (track/loc (heap-ref (+ loc x)))))
        (heap-set! loc 'frwd)
        (heap-set! (+ loc 1) new-addr)
        new-addr]
       [(frwd) (heap-ref (+ loc 1))]
       [else (error 'move/loc "wrong tag ~s at ~a" (heap-ref loc) loc)])]
    [else loc]))

;; move/ref : loc -> loc
;; move the referenced object to the other semi-space
;; and return the new addr of moved object
(define (move/ref loc)
  (unless (2nd-gen? loc)
    (error 'move/loc "wrong location ~s" loc))
  (case (heap-ref loc)
    [(flat grey-flat white-flat) (void)]
    [(pair grey-pair white-pair)
     (gc:set-first! loc (move/loc (heap-ref (+ loc 1))))
     (gc:set-rest! loc (move/loc (heap-ref (+ loc 2))))
     (move/ref (heap-ref (+ loc 1)))
     (move/ref (heap-ref (+ loc 2)))]
    [(proc grey-proc white-proc)
     (define fv-count (heap-ref (+ loc 2)))
     (for ([x (in-range 0 fv-count)])
       (define l (+ loc 3 x))
       (heap-set! l (move/loc (heap-ref l)))
       (move/ref (heap-ref l)))]
    [else (error 'move/ref "wrong tag at ~a" loc)]))

(define (move/pointers loc)
  (cond
    [(or (= loc (2nd-gen-start))
         (equal? 'free (heap-ref loc)))
     (heap-set! (table-start-word) (add1 (table-start-word)))]
    [else
     (define new-addr (move/loc (heap-ref (+ loc 1))))
     (heap-set! (heap-ref loc) new-addr)
     (move/ref new-addr)
     (heap-set! loc 'free)
     (heap-set! (+ loc 1) 'free)
     (move/pointers (+ loc 2))]))

(define (free-1st-gen)
  (for ([i (in-range 1 (1st-gen-size))])
    (heap-set! i 'free)))

(define (1st-gen? loc)
  (and (>= loc 1)
       (< loc (1st-gen-size))))

(define (2nd-gen? loc)
  (and (>= loc (2nd-gen-alloc-start))
       (< loc (2nd-gen-end))))

;; find-free-space : find free space by traversing free space list
;; layout := free-2 next
;;        |  free-n next size
;; next := (or/c location? #f)
;; free-list-head : holds the value of head of free space list
;;                | #f means already runs out of free space
;;
;; before give a series of free spaces to alloc function
;; must update the free space list
;; iff its the first free space in list must also udpate free-list-head as well
(define (find-free-space start prev size)
  (local [(define (next-in-free-list loc)
            (heap-ref (+ loc 1)))
          (define (update-next-in-prev prev loc)
            (heap-set! (if prev
                           (+ prev 1)
                           (free-list-head))
                       loc))]
    (cond
      [(not start) #f]
      [else
       (case (heap-ref start)
         [(free-2)
          (cond
            [(= size 2)
             (update-next-in-prev prev (next-in-free-list start))
             start]
            [else (find-free-space (heap-ref (+ start 1)) start size)])]
         [(free-n)
          (define length (heap-ref (+ start 2)))
          (cond
            [(= size length)
             (update-next-in-prev prev (next-in-free-list start))
             start]
            [(< size length)
             (define new-free (+ start size))
             (define new-size (- length size))
             (cond
               [(= new-size 1)
                (update-next-in-prev prev (next-in-free-list start))
                (heap-set! new-free 'free)
                start]
               [else
                (update-next-in-prev prev new-free)
                (cond
                  [(= new-size 2)
                   (heap-set! new-free 'free-2)
                   (heap-set! (+ new-free 1) (heap-ref (+ start 1)))]
                  [else (heap-set! new-free 'free-n)
                        (heap-set! (+ new-free 1) (heap-ref (+ start 1)))
                        (heap-set! (+ new-free 2) new-size)])
                start])]
            [else (find-free-space (heap-ref (+ start 1)) start size)])]
         [else (error 'find-free-space "wrong tag @ ~s" start)])])))

(define (make-pointers-to-2nd-gen-roots start)
  (cond
    [(= start (1st-gen-size)) (void)]
    [else
     (case (heap-ref start)
       [(flat) (make-pointers-to-2nd-gen-roots (+ start 2))]
       [(pair) (define one-loc (heap-ref (+ start 1)))
               (when (2nd-gen? one-loc) (trace/roots-iff-white one-loc))
               (define another-loc (heap-ref (+ start 2)))
               (when (2nd-gen? another-loc) (trace/roots-iff-white another-loc))
               (make-pointers-to-2nd-gen-roots (+ start 3))]
       [(proc) (define fv-counts (heap-ref (+ start 2)))
               (for ([i (in-range fv-counts)])
                 (define loc (heap-ref (+ start 3 i)))
                 (when (2nd-gen? loc) (trace/roots-iff-white loc)))
               (make-pointers-to-2nd-gen-roots (+ start 3 fv-counts))]
       [(frwd) (define loc (heap-ref (+ start 1)))
               (traverse/roots loc)
               (case (heap-ref loc)
                 [(flat) (make-pointers-to-2nd-gen-roots (+ start 2))]
                 [(pair) (make-pointers-to-2nd-gen-roots (+ start 3))]
                 [(proc) (make-pointers-to-2nd-gen-roots (+ start 3 (heap-ref (+ loc 2))))])]
       [(free) (void)]
       [else (error 'make-pointers-to-2nd-gen-roots "wrong tag at ~a" start)])]))

;; check if any spot within free slots are taken
;; in order to detect heap-stack crash
(define (check/free loc size)
  (case (heap-ref loc)
    [(free-2) #t]
    [(free-n) (stream-andmap (lambda (x)
                               (equal? 'free (heap-ref (+ loc x))))
                             (in-range 3 size))]
    [else (error 'check/free "wrong tag @ ~s" loc)]))

(define (copy/alloc n some-roots more-roots)
  (define next (find-free-space (heap-ref (free-list-head)) #f n))
  (unless next (error 'copy/alloc "out of memory"))
  (unless (check/free next n)
    (error 'copy/alloc "heap hit tracing stack at ~s" next))
  ;; incremental collection
  (heap-set!
   (step-count-word)
   (+ n (heap-ref (step-count-word))))
  (when (>= (heap-ref (step-count-word))
            step-length)
    (traverse/incre-mark (next/cont)))
  next)

(define (trace/roots-iff-white thing)
  (cond
    [(list? thing)
     (for-each trace/roots-iff-white thing)]
    [(root? thing)
     (define loc (read-root thing))
     (when (white? loc)
       (push/cont loc))]
    [(number? thing)
     (when (white? thing)
       (push/cont thing))]))

;; object-length : location -> number
(define (object-length loc)
  (define tag (heap-ref loc))
  (case tag
    [(free) 1]
    [(free-2) 2]
    [(free-n) (heap-ref (+ loc 2))]
    [(flat white-flat) 2]
    [(pair white-pair) 3]
    [(proc white-proc) (+ 3 (heap-ref (+ loc 2)))]
    [else (error 'object-length "wrong tag ~s @ ~s" tag loc)]))

(define (mark-white! i)
  (case (heap-ref i)
    [(flat) (heap-set! i 'white-flat)]
    [(pair) (heap-set! i 'white-pair)]
    [(proc) (heap-set! i 'white-proc)]
    [else (error 'mark-white! "wrong tag at ~a" i)]))

;; for free-n spaces, free its rest slots besides the first three
;; free-rest : location -> void
(define (free-rest loc)
  (unless (equal? 'free-n (heap-ref loc))
    (error 'free-rest "wrong tag @ ~s" loc))

  (for ([i (in-range 3 (heap-ref (+ loc 2)))])
    #:break (equal? 'free (heap-ref (+ loc i)))
    (heap-set! (+ loc i) 'free)))

(define (free/mark-white! loc prev last-start so-far)
  (unless (or (and last-start so-far)
              (not (or last-start so-far)))
    (error 'free/mark-white!
           "cumulating info are incorrect, last-start: ~s, so-far: ~s"
           last-start so-far)) ;⋮
  (cond
    [(>= loc (2nd-gen-end))
     (cond
       [(and last-start so-far)
        (cond
          [(= 1 so-far) (void)]
          [(= 2 so-far) (heap-set! last-start 'free-2)
                               (heap-set! (+ 1 last-start) #f)
                               (heap-set! (if prev (+ prev 1) (free-list-head))
                                             last-start)]
          [else (heap-set! last-start 'free-n)
                (heap-set! (+ 1 last-start) #f)
                (heap-set! (+ 2 last-start) so-far)
                (heap-set! (if prev (+ prev 1) (free-list-head))
                              last-start)
                (free-rest last-start)])]
       [else (void)])#;(⋮)]
    [else
     (case (heap-ref loc)
       [(flat pair proc) (mark-white! loc)
        (define length (object-length loc))
        (cond
          [(and last-start
                so-far
                (= 1 so-far))
           (free/mark-white! (+ loc length) prev #f #f)]
          [(and last-start
                so-far
                (>= so-far 2))
           (cond
             [(= 2 so-far) (heap-set! last-start 'free-2)
                                  (heap-set! (+ last-start 1) #f)]
             [else (heap-set! last-start 'free-n)
                   (heap-set! (+ last-start 1) #f)
                   (heap-set! (+ last-start 2) so-far)])
           (if prev
               (heap-set! (+ prev 1) last-start)
               (heap-set! (free-list-head) last-start))
           (free/mark-white! (+ loc length) last-start #f #f)]
          [else (free/mark-white! (+ loc length) prev #f #f)])#;(⋮)]
       [(white-flat white-pair white-proc
                    free free-2 free-n)
        (let* ([length (object-length loc)]
               [join (and last-start so-far)]
               [start (if join last-start loc)]
               [newlen (if join (+ so-far length) length)])
          (free/mark-white!
           (+ loc length) prev start newlen))]
       [else (error 'free/mark-white!
                    "wrong tag at ~a" loc)])]))

(define (traverse/incre-mark loc)
  (cond
    [(not loc) (heap-set! (step-count-word) 0)]
    [else
     (case (heap-ref loc)
       [(flat grey-flat)
        (mark-black loc)
        (step/count 2)
        (continue/incre-mark)] ; ⋮
       [(pair grey-pair)
        (mark-black loc)
        (step/count 3) ; 3 words scanned
        (push/cont (heap-ref (+ loc 2)))
        (push/cont (heap-ref (+ loc 1)))
        (continue/incre-mark)]
       [(proc grey-proc)
        (mark-black loc)
        (define closure-size (heap-ref (+ loc 2)))
        (step/count (+ 3 closure-size))
        (for ([i (in-range closure-size)])
          (push/cont (heap-ref (+ loc 3 i))))
        (continue/incre-mark)] ; ⋮
       [else (error 'traverse/incre-mark
                    "wrong tag at ~s" loc)])]))

(define (mark-black loc)
  (case (heap-ref loc)
    [(grey-flat)
     (heap-set! loc 'flat)]
    [(grey-pair)
     (mark-grey (heap-ref (+ loc 1)))
     (mark-grey (heap-ref (+ loc 2)))
     (heap-set! loc 'pair)]
    [(grey-proc)
     (for ([i (in-range (heap-ref (+ loc 2)))])
       (mark-grey (heap-ref (+ loc 3 i))))
     (heap-set! loc 'proc)]
    [(pair flat proc white-pair white-flat white-proc  cont) (void)]))

(define (mark-grey loc)
  (case (heap-ref loc)
    [(white-pair) (heap-set! loc 'grey-pair)]
    [(white-flat) (heap-set! loc 'grey-flat)]
    [(white-proc) (heap-set! loc 'grey-proc)]
    [(pair flat proc grey-pair grey-flat grey-proc cont) (void)]))

(define (heap-cont/check)
  (check/tag (2nd-gen-alloc-start))
  (heap-check (2nd-gen-alloc-start)))

(define (heap-check loc)
  (when (< loc (2nd-gen-end))
    (case (heap-ref loc)
      [(pair)
       (if (or (white? (heap-ref (+ loc 1)))
               (white? (heap-ref (+ loc 2))))
           (error 'heap-check "black object points to white object at ~a" loc)
           (heap-check (+ loc 3)))]
      [(flat)
       (heap-check (+ loc 2))]
      [(proc)
       (define closure-size (heap-ref (+ loc 2)))
       (case closure-size
         [(free)
          (heap-check (+ loc 3))]
         [else
          (for ([i (in-range closure-size)])
            (when (white? (heap-ref (+ loc 3 i)))
              (error 'heap-check "black object points to white object at ~a" loc)))
          (heap-check (+ loc 3 closure-size))])]
      [(white-pair grey-pair)
       (heap-check (+ loc 3))]
      [(white-flat grey-flat)
       (heap-check (+ loc 2))]
      [(white-proc grey-proc)
       (heap-check (+ loc 3 (heap-ref (+ loc 2))))]
      [(free) (heap-check (+ loc 1))]
      [else
       (cond
         [(location? loc) (check/tag-helper loc)]
         [else (error 'check "unknown tag @ ~a" loc)])])))

(define (check/tag loc)
  (when (< loc (2nd-gen-end))
    (case (heap-ref loc)
      [(pair white-pair grey-pair) (check/tag (+ loc 3))]
      [(flat white-flat grey-flat) (check/tag (+ loc 2))]
      [(proc white-proc grey-proc)
       (define closure-size (heap-ref (+ loc 2)))
       (check/tag (+ loc 3 closure-size))]
      [(free) (check/tag (+ loc 1))]
      [else
       (cond
         [(location? loc) (check/tag-helper loc)]
         [else (error 'check/tag "wrong tag at ~a" loc)])])))

(define (check/tag-helper loc)
  (for ([x (in-range loc (2nd-gen-end))])
    (unless (location? (heap-ref x))
      (error 'check/tag-helper "wrong value @ ~s, should be location" x))))

;; white? : location? -> boolean?
(define (white? loc)
  (case (heap-ref loc)
    [(white-pair white-proc white-flat)
     true]
    [else false]))

(define (next/cont)
  (define loc (heap-ref (tracing-head-word)))
  (cond
    [(equal? loc #f) #f]
    [else
     (define ptr (heap-ref loc))
     (clean/cont loc)
     ptr]))

(define (push/cont ptr)
  (mark-grey ptr)
  (when (not (member (heap-ref ptr) '(pair flat proc)))
    (define loc (heap-ref (tracing-head-word)))
     (when (and loc
                (not (equal? 'free (heap-ref (sub1 loc)))))
       (error 'push/cont "collection crashed at ~s" (sub1 loc))) ; ⋮
     (let ([top
            (if loc (sub1 loc) (- (2nd-gen-end) 1))])
       (heap-set! top ptr)
       (heap-set! (tracing-head-word) top))))

(define (continue/incre-mark)
  (if (step/finished?)
      (heap-set! (step-count-word) 0)
      (traverse/incre-mark (next/cont))))

(define (clean/cont loc)
  (define next (+ loc 1))
  (when (> next (2nd-gen-end))
    (error 'clean/cont "stack is out of bound"))

  (cond
    [(= next (2nd-gen-end)) (heap-set! (tracing-head-word) #f)]
    [else (heap-set! (tracing-head-word) (+ loc 1))])
  (heap-set! loc 'free))

(define (step/count n)
  (heap-set! (step-count-word) (- (heap-ref (step-count-word)) n)))

(define (step/finished?)
  (<= (heap-ref (step-count-word)) 0))

; validate-heap : -> void?
(define (validate-heap from to)
  (define (valid-pointer? p)
    (unless (< p (heap-size))
      (error 'validate-heap "pointer out of bounds ~a" p))
    (unless (member (heap-ref p) '(flat pair proc
                                        grey-flat grey-pair grey-proc
                                        white-flat white-pair white-proc))
      (error 'validate-heap "pointer to non-tag ~a" p)))
  (let loop ([i from])
    (when (< i to)
      (case (heap-ref i)
        [(white-flat grey-flat flat) (loop (+ i 2))]
        [(white-pair grey-pair pair)
         (valid-pointer? (heap-ref (+ i 1)))
         (valid-pointer? (heap-ref (+ i 2)))
         (loop (+ i 3))]
        [(white-proc grey-proc proc)
         (for ([j (in-range 0 (heap-ref (+ i 2)))])
           (valid-pointer? (heap-ref (+ i 3 j))))
         (loop (+ i 3 (heap-ref (+ i 2))))] ; use stored length
        [(free-n)
         (define next (heap-ref (+ i 1)))
         (or (not next) (valid-pointer? next))
         (define len (heap-ref (+ i 2)))
         (unless (and (number? len)
                      (> len 2)
                      (<= len (- (2nd-gen-end) (2nd-gen-alloc-start))))
           (error 'validate-heap "bad free list node at ~a" i))
         (loop (+ i len))]
        [(free-2) (loop (+ i 2))]
        [(free) (loop (+ i 1))]
        [else (error 'validate-heap
                     "unexpected tag ~a at ~a" (heap-ref i) i)]))); @$\vdots$@
  )

(define (validate-nursery)
  (let ([end (heap-ref 0)])
    (validate-heap 1 end)))

(define (validate-2nd-gen)
  (let ([start (2nd-gen-alloc-start)]
        [end (2nd-gen-end)])
    (validate-heap start end)))

(define (validate-both-heaps)
  (validate-nursery)
  (validate-2nd-gen)
  #t)

(module+ test
  (print-only-errors)
  (with-heap (make-vector 256)
    (init-allocator)
    (test (validate-both-heaps) #t)
    ;; fill nursery, but don't force gc
    (for ([i (in-range 31)])
      (gc:alloc-flat i))
    (test (validate-nursery) (void))
    (test (validate-2nd-gen) (void))
    (test (vector-copy (current-heap) 0 3) #(63 flat 0))
    (test (vector-copy (current-heap) 100 103) #(free-n #f 156))
    ;; force minor collection, start at beginning
    (define a-root 1)
    (with-roots (a-root)
      (gc:alloc-flat 32))
    (test (validate-both-heaps) #t)
    (test (vector-copy (current-heap) 0 3) #(3 flat 32))
    (test (vector-copy (current-heap) 100 105) #(white-flat 0 free-n #f 154))
    )
  ;; allocation tests with cons pairs
  (with-heap (make-vector 256)
    (init-allocator)
    (test (validate-nursery) (void))
    ;; fill nursery, but don't force gc
    (for ([i (in-range 9)])
      (gc:cons
       (simple-root (gc:alloc-flat (* 2 i)))
       (simple-root (gc:alloc-flat (add1 (* 2 i))))))
    (test (validate-both-heaps) #t)
    (test (vector-copy (current-heap) 0 8)  #(64 flat 0 flat 1 pair 1 3))
    (test (vector-copy (current-heap) 100 103) #(free-n #f 156))
    ;; force minor collection, start at beginning
    (define a-root 5)
    (with-roots (a-root)
      (gc:alloc-flat 32))
    (test (validate-both-heaps) #t)
    (test (vector-copy (current-heap) 0 3) #(3 flat 32))
    (test (vector-copy (current-heap) 100 107) #(white-pair 103 105 white-flat 0 white-flat 1)))

  ;; one long list
  (with-heap (make-vector 256)
    (init-allocator)
    (test (validate-both-heaps) #t)
    ;; fill nursery, but don't force gc
    (define head
      (for/fold ([lst (gc:alloc-flat empty)])
                ([i (in-range 12)])
        (gc:cons
         (simple-root (gc:alloc-flat i))
         (simple-root lst))))
    (test (validate-both-heaps) #t)
    (with-roots (head)
      (gc:alloc-flat 13))
    (test (validate-both-heaps) #t)
    (test head 100)
    (test (vector-copy (current-heap) 100 105) #(white-pair 103 105 white-flat 11))
    )

  ;; mostly fill memory with a list
  (with-heap (make-vector 256)
     (init-allocator)
     (define head
       (for/fold ([lst (gc:alloc-flat empty)])
                 ([i (in-range 37)])
         (validate-both-heaps)
         (with-roots (lst)
           (gc:cons
            (simple-root (gc:alloc-flat i))
            (simple-root lst)))))
    (test/exn 
     (with-roots (head)
       (gc:alloc-flat 'force-collection))
     "out of memory")
    )

(with-heap (make-vector 256)
  (init-allocator)
  (define head
    (for/fold ([lst (gc:alloc-flat empty)])
              ([i (in-range 30)])
      (validate-both-heaps)
      (with-roots (lst)
         (gc:cons
          (simple-root (gc:alloc-flat i))
          (simple-root lst)))))
  (validate-2nd-gen)
  (define head2
    (for/fold ([lst (gc:alloc-flat empty)])
               ([i (in-range 31 37)])
       (validate-both-heaps)
       (with-roots (lst head)
         (gc:cons
          (simple-root (gc:alloc-flat i))
          (simple-root lst)))))
  ;; force nursery collection
  (for ([i (in-range 33)])
    (with-roots (head2)
      (gc:alloc-flat i)))
  (validate-2nd-gen)
  )
)