#lang racket
(require [only-in plait test test/exn error print-only-errors])

;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;
;; This part provides a small part of the plai/gc2/collector language
(define current-heap (make-parameter (make-vector 0 #f)))
(define (heap-set! index val) (vector-set!  (current-heap) index val))
(define (heap-ref index) (vector-ref (current-heap) index))
(define (heap-size) (vector-length (current-heap)))
(define-syntax-rule (with-heap vec expr ...)
  (parameterize
      ([current-heap vec])
    (begin
      expr ...)))

;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;
;; This part is the partial implementation of a collector API For
;; simplicity, use our very first allocator without any data
;; structure.

(define (init-allocator)
  (vector-fill! (current-heap) 'free))

(define (alloc/header tag . vals)
  (define loc (malloc (+ 1 (length vals))))
  (heap-set! loc tag)
  (for ([i (in-range (length vals))]
        [v (in-list vals)])
    (heap-set! (+ loc i 1) v))
  loc)

(define (gc:alloc-flat val) (alloc/header 'flat val))

(define (gc:cons val1 val2) (alloc/header 'cons val1 val2))


;; Linear time allocator, based on Lecture 20

(define (malloc size)
  (define ptr (find-free-space size))
  (unless ptr  (error 'alloc "out of memory"))
  ptr)

(module+ test
  (with-heap (make-vector 10 #f)
    (init-allocator)
    (test/exn (malloc 100) "out of memory")))

(define (find-free-space n)
  (define (n-free-blocks? start n)
    (for/fold ([ok #t])
              ([i (in-range start (+ start n))])
      (and ok (< i (heap-size)) (equal? (heap-ref i) 'free))))

  (define (loop start)
    (and
     (< start (heap-size))
     (case (heap-ref start)
       [(flat) (loop (+ start 2))]
       [(cons) (loop (+ start 3))]
       [(free) (if (n-free-blocks? start n)
                   start
                   (loop (+ start 1)))]
       [else (error 'find-free-space
                    "unexpected tag ~a" start)])))
  (loop 0))

;; individual passes, from lecture 20
#;(define (mark-white!)
    (let loop ([i 0])
      (when (< i (heap-size))
        (case (heap-ref i)
          [(cons)
           (heap-set! i 'white-cons) (loop (+ i 3))]
          [(flat)
           (heap-set! i 'white-flat) (loop (+ i 2))]
          [(free) (loop (+ i 1))]
          [else (error 'mark-white! "bad tag: ~a" i)]))))

#;(define (free-white!)
  (let loop ([i 0])
    (when (< i (heap-size))
      (case (heap-ref i)
        [(cons) (loop (+ i 3))]
        [(flat) (loop (+ i 2))]
        [(free) (loop (+ i 1))]
        [(white-flat) (heap-set! i 'free)
                      (heap-set! (+ i 1) 'free)
                      (loop (+ i 2))]
        [(white-cons) (heap-set! i 'free)
                      (heap-set! (+ i 1) 'free)
                      (heap-set! (+ i 2) 'free)
                      (loop (+ i 3))]
        [else (error 'free-white! "unexpected tag: ~a" i)]))))

;; Here is the function you need to write 
(define (free/mark-white!)
  (void))

(module+ test
  (with-heap (vector 'free 'flat 1)
    (free/mark-white!)
    (test (current-heap) #(free white-flat 1)))

  (with-heap (vector 'free 'white-flat 1)
    (free/mark-white!)
    (test (current-heap) #(free free free)))

  (with-heap (vector 'free 'cons 1 1)
    (free/mark-white!)
    (test (current-heap) #(free white-cons 1 1)))

  (with-heap (vector 'free 'cons 1 1)
    (free/mark-white!)
    (free/mark-white!)
    (test (current-heap) #(free free free free)))

  (with-heap (vector 'white-flat 0 'cons 0 0)
    (free/mark-white!)
    (test (current-heap) #(free free white-cons 0 0)))

  (with-heap (make-vector 7 '?)
    (init-allocator)
    (test (current-heap) (make-vector 7 'free))
    (gc:alloc-flat 'first)
    (test (current-heap) #(flat first free free free free free))
    (gc:alloc-flat 'rest)
    (test (current-heap) #(flat first flat rest free free free))
    (gc:cons 0 2)
    (test (current-heap) #(flat first flat rest cons 0 2))
    (free/mark-white!)
    (test (current-heap) #(white-flat first white-flat rest white-cons 0 2))
    (free/mark-white!)
    (test (current-heap)  (make-vector 7 'free))
    (gc:alloc-flat 'first)
    (gc:alloc-flat 'rest)
    (gc:cons 0 2)
    (free/mark-white!)
    (test (current-heap) #(white-flat first white-flat rest white-cons 0 2))
    )

  (with-heap  (vector 'flat 'first 'flat 'rest 'white-cons 0 2)
    (free/mark-white!)
    (test (current-heap) #(white-flat first white-flat rest free free free))
    (heap-set! 0 'flat)
    (free/mark-white!)
    (test (current-heap) #(white-flat first free free free free free)))

  (with-heap (make-vector 5 #f)
    (init-allocator)
    (gc:cons 0 0)
    (test (current-heap) #(cons 0 0 free free))
    (malloc 2)
    (test/exn (malloc 100) "out of memory")
    (heap-set! 0 'fail)
    (test/exn (malloc 2) "unexpected tag"))
  )