(import (chezscheme)
        (bcache block-manager)
	(fmt fmt)
	(loops)
        (persistent-data btree)
	(thin metadata))

(define (noop k v)
  'nil)

(define (mapping-top-level root)
  (btree-open le64-vt root))

(define (mapping-bottom-level root)
  (btree-open le64-vt root))

(define (mk-incrementer counts)
 (lambda (i)
  (vector-set! counts i (+ (vector-ref counts i) 1))))

(define (bottom-level-walker cache inc-fn)
 (lambda (key root)
  (btree-each-and-count (mapping-bottom-level root) cache noop inc-fn)))

(define (walk-top-level cache root counts)
  (let ((inc-fn (mk-incrementer counts)))
   (btree-each-and-count (mapping-top-level root)
    			 cache
			 (bottom-level-walker cache inc-fn)
			 inc-fn)))

(define (calc-ref-counts md-path)
  (with-bcache (cache md-path (* 1 1024))
    (let ((counts (make-vector (get-nr-blocks cache) 0)))
     (with-block (b cache 0 (get-flags))
      (let ((sb (block->superblock b)))
	(walk-top-level cache
	                (ftype-ref ThinSuperblock (data-mapping-root) sb)
			counts)))

     counts)))

(define (print-counts counts)
  (upto (i (vector-length counts))
   (let ((rc (vector-ref counts i)))
    (when (> rc 0)
     (fmt #t i ": " (vector-ref counts i) nl)))))

(define (usage)
 (fmt #t "Usage: thin-calc-ref-counts <binary metadata>" nl)
 (exit 1))

(let ((args (cdr (command-line))))
  (if (not (= (length args) 1))
   (usage)
   (let ((counts (calc-ref-counts (car args))))
    (print-counts counts))))

