;; Copyright (C) 2015 Free Software Foundation, Inc.

;;;; This library is free software; you can redistribute it and/or
;;;; modify it under the terms of the GNU Lesser General Public
;;;; License as published by the Free Software Foundation; either
;;;; version 3 of the License, or (at your option) any later version.
;;;;
;;;; This library is distributed in the hope that it will be useful,
;;;; but WITHOUT ANY WARRANTY; without even the implied warranty of
;;;; MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU
;;;; Lesser General Public License for more details.
;;;;
;;;; You should have received a copy of the GNU Lesser General Public
;;;; License along with this library; if not, write to the Free Software
;;;; Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301 USA

;;; Commentary:
;;;
;;; A fector is a functional vector.
;;;
;;; Code:

(define-module (fector)
  #:use-module (srfi srfi-9)
  #:use-module (ice-9 match)
  #:export (fector?
	    transient-fector?
	    transient-fector
	    persistent-fector
	    empty-fector
	    make-fector
	    fector-length
	    fector-ref
	    fector-set
	    fector-set!
	    fector-push
	    fector-push!
	    fector-pop
	    fector-pop!
	    fector-fold
	    fector-fold-right))

(define-syntax-rule (define-inline name val)
  (define-syntax name (identifier-syntax val)))

;; FIXME: This should make an actual atomic reference.
(define-inlinable (make-atomic-reference value)
  (list value))
(define-inlinable (get-atomic-reference reference)
  (car reference))
(define-inlinable (set-atomic-reference! reference value)
  (set-car! reference value))

(define-inline *branch-bits* 5)
(define-inline *branch-size* (ash 1 *branch-bits*))
(define-inline *branch-size-with-edit* (1+ *branch-size*))
(define-inline *edit-index* *branch-size*)
(define-inline *branch-mask* (1- *branch-size*))

(define-record-type <fector>
  (make-fector length shift root tail)
  fector?
  (length %fector-length)
  (shift fector-shift)
  (root fector-root)
  (tail fector-tail))

(define-record-type <transient-fector>
  (make-transient-fector length shift root tail edit)
  transient-fector?
  (length transient-fector-length set-transient-fector-length!)
  (shift transient-fector-shift set-transient-fector-shift!)
  (root transient-fector-root set-transient-fector-root!)
  (tail transient-fector-tail set-transient-fector-tail!)
  (edit transient-fector-edit set-transient-fector-edit!))

(define (fector-length fector)
  (match fector
    (($ <fector> length) length)
    (($ <transient-fector> length) length)))

(define-inlinable (new-branch edit)
  (let ((vec (make-vector *branch-size-with-edit* #f)))
    (when edit (vector-set! vec *edit-index* edit))
    vec))
(define (clone-branch-and-set branch i elt)
  (let ((new (vector-copy branch)))
    (vector-set! new i elt)
    new))
(define-inlinable (assert-readable! root-edit)
  (unless (eq? (get-atomic-reference root-edit) (current-thread))
    (error "Transient fector owned by another thread" root-edit)))
(define-inlinable (writable-branch branch root-edit)
  (let ((edit (vector-ref branch *edit-index*)))
    (if (eq? root-edit edit)
        branch
        (clone-branch-and-set branch *edit-index* root-edit))))

(define-inlinable (round-down min shift)
  (logand min (lognot (1- (ash 1 shift)))))

(define-inlinable (compute-tail-offset length)
  (if (< length *branch-size*)
      0
      (logand (1- length) (lognot *branch-mask*))))

(define-inlinable (compute-tail-length length)
  (- length (compute-tail-offset length)))

(define *empty-branch* (new-branch #f))

(define empty-fector (make-fector 0 *branch-bits* *empty-branch* #()))

(define* (transient-fector #:optional (source empty-fector))
  (match source
    (($ <transient-fector> length shift root tail edit)
     (assert-readable! edit)
     source)
    (($ <fector> length shift root tail)
     (let* ((edit (make-atomic-reference (current-thread)))
	    (tail* (new-branch edit)))
       (vector-move-left! tail 0 (vector-length tail) tail* 0)
       (make-transient-fector length shift root tail* edit)))))

(define* (persistent-fector #:optional (source empty-fector))
  (match source
    (($ <transient-fector> length shift root tail edit)
     (assert-readable! edit)
     ;; Make a fresh reference, causing any further operations on this
     ;; transient to clone its root afresh.
     (set-transient-fector-edit! source
                                 (make-atomic-reference (current-thread)))
     ;; Clear the reference to the current thread, causing our edited
     ;; data structures to be persistent again.
     (set-atomic-reference! edit #f)
     (let* ((tail-length (compute-tail-length length))
	    (tail* (make-vector tail-length #f)))
       (vector-move-left! tail 0 tail-length tail* 0)
       (make-fector length shift root tail)))
    (($ <fector>)
     source)))

(define (fector-set! fector i val)
  (define (update-tree! shift root)
    (let* ((shift (- shift *branch-bits*))
	   (idx (logand (ash i (- shift)) *branch-mask*)))
      (if (zero? shift)
	  (vector-set! root idx val)
	  (let* ((edit (vector-ref root *edit-index*))
		 (branch (vector-ref root idx))
		 (branch* (writable-branch branch edit)))
	    (unless (eq? branch branch*)
	      (vector-set! root idx branch*))
	    (update-tree! shift branch*)))))
  (match fector
    (($ <transient-fector> length shift root tail edit)
     (assert-readable! edit)
     (unless (and (< 0 i) (< i length))
       (error "index out of range" i))
     (let ((tail-offset (compute-tail-offset length)))
       (cond
	((<= tail-offset i)
	 (vector-set! tail (- i tail-offset) val))
	(else
	 (let ((root* (writable-branch root edit)))
	   (unless (eq? root root*)
	     (set-transient-fector-root! fector root*))
	   (update-tree! shift root*)))))
     fector)
    (($ <fector>)
     (fector-set! (transient-fector fector) i val))))

(define (fector-set fector i val)
  (define (update-tree shift root)
    (let* ((shift (- shift *branch-bits*))
	   (idx (logand (ash i (- shift)) *branch-mask*))
	   (node (if (zero? shift)
		     val
		     (update-tree shift (vector-ref root idx)))))
      (clone-branch-and-set root idx node)))
  (match fector
    (($ <fector> length shift root tail)
     (unless (and (< 0 i) (< i length))
       (error "index out of range" i))
     (let ((tail-offset (compute-tail-offset length)))
       (cond
	((<= tail-offset i)
	 (let ((new-tail (vector-copy tail)))
	   (vector-set! new-tail (- i tail-offset) val)
	   (make-fector length shift root new-tail)))
	(else
	 (make-fector length shift (update-tree shift root) tail)))))
    (($ <transient-fector>)
     (fector-set (persistent-fector fector) i val))))

(define (fector-push! fector val)
  (define (add-tail! i shift root tail)
    (let* ((shift (- shift *branch-bits*))
	   (idx (logand (ash i (- shift)) *branch-mask*)))
      (if (= shift *branch-bits*)
	  (vector-set! root idx tail)
	  (let* ((branch (vector-ref root idx))
		 (edit (vector-ref root *edit-index*))
		 (branch* (if branch
			      (writable-branch branch edit)
			      (new-branch edit))))
	    (unless (eq? branch branch*)
	      (vector-set! root idx branch*))
	    (add-tail! i shift branch* tail)))))
  (match fector
    (($ <transient-fector> length shift root tail edit)
     (assert-readable! edit)
     (let* ((tail-length (compute-tail-length length)))
       (cond
	((< tail-length *branch-size*)
	 ;; Normal case: just add to the tail.
	 (vector-set! tail tail-length val))
	(else
	 ;; Tail array is full; push into tree.
	 (cond
	  ((= length *branch-size*)
	   ;; The tail becomes the first root.
	   (set-transient-fector-root! fector tail))
	  ((= length (+ (ash 1 shift) *branch-size*))
	   ;; Tree is full; add a level.
	   (let ((root* (new-branch edit))
		 (shift* (+ shift *branch-bits*)))
	     (vector-set! root* 0 root)
	     (set-transient-fector-root! fector root*)
	     (set-transient-fector-shift! fector shift*)
	     (add-tail! (- length *branch-size*) shift* root* tail)))
	  (else
	   (let ((root* (writable-branch root edit)))
	     (unless (eq? root root*)
	       (set-transient-fector-root! fector root*))
	     (add-tail! (- length *branch-size*) shift root* tail))))
	 ;; Make a fresh tail and add the pushed val.
	 (let ((tail (new-branch edit)))
	   (set-transient-fector-tail! fector tail)
	   (vector-set! tail 0 val)))))
     (set-transient-fector-length! fector (1+ length))
     fector)
    (($ <fector>)
     (fector-push! (transient-fector fector) val))))

(define (fector-push fector val)
  (define (add-tail i shift root tail)
    (let* ((shift (- shift *branch-bits*))
	   (idx (logand (ash i (- shift)) *branch-mask*))
	   (branch (if (= shift *branch-bits*)
		       tail
		       (add-tail i shift
				 (or (vector-ref root idx) (new-branch #f))
				 tail))))
      (clone-branch-and-set root idx branch)))
  (match fector
    (($ <fector> length shift root tail)
     (let* ((tail-length (compute-tail-length length)))
       (cond
	((< tail-length *branch-size*)
	 ;; Normal case: just add to the tail.
	 (let ((new-tail (make-vector (1+ tail-length) #f)))
	   (vector-move-left! tail 0 tail-length new-tail 0)
	   (vector-set! new-tail tail-length val)
	   (make-fector (1+ length) shift root new-tail)))
	(else
	 ;; Tail array is full; push into tree.  We have to copy the
	 ;; tail, though, in order to add the trailing "edit" field.
	 (let ((tail (let ((tail* (new-branch #f)))
		       (vector-move-left! tail 0 *branch-size* tail* 0)
		       tail*))
	       (new-tail (new-branch #f)))
	   ;; Go ahead and add the pushed value to the new tail.
	   (vector-set! new-tail 0 val)
	   (cond
	    ((= length *branch-size*)
	     ;; The tail becomes the first root.
	     (make-fector (1+ length) shift tail new-tail))
	    ((= length (+ (ash 1 shift) *branch-size*))
	     ;; Tree is full; add a level.
	     (let ((root* (new-branch #f))
		   (shift* (+ shift *branch-bits*)))
	       (vector-set! root* 0 root)
	       (let ((root* (add-tail (ash 1 shift) shift* root* tail)))
		 (make-fector (1+ length) shift* root* new-tail))))
	    (else
	     (let ((root (add-tail (- length *branch-size*) shift root tail)))
	       (make-fector (1+ length) shift root new-tail)))))))))
    (($ <transient-fector>)
     (fector-push (persistent-fector fector) val))))

(define (fector-pop! fector)
  (define (pop-tail! i shift root)
    (let* ((shift (- shift *branch-bits*))
	   (idx (logand (ash i (- shift)) *branch-mask*)))
      (if (= shift *branch-bits*)
	  (let ((tail (vector-ref root idx)))
	    (vector-set! root idx #f)
	    tail)
	  (let* ((branch (vector-ref root idx))
		 (edit (vector-ref root *edit-index*))
		 (branch* (writable-branch branch edit)))
	    (unless (eq? branch branch*)
	      (vector-set! root idx branch*))
	    (pop-tail! i shift branch*)))))
  (match fector
    (($ <transient-fector> length shift root tail edit)
     (assert-readable! edit)
     (let* ((tail-length (compute-tail-length length)))
       (cond
	((< 1 tail-length)
	 ;; Normal case: just clear the last entry in the tail.
	 (when (zero? length)
	   (error "can't pop from empty fector"))
	 (vector-set! tail (1- tail-length) #f))
	(else
	 ;; Tail array will be empty; pop a tail from the tree.
	 (cond
	  ((= length (1+ *branch-size*))
	   ;; The root becomes the tail.
	   (set-transient-fector-tail! fector (writable-branch root edit))
	   (set-transient-fector-root! fector *empty-branch*))
	  ((= length (+ (ash 1 shift) *branch-size* 1))
	   ;; Shrink the tree.
	   (let ((tail (let lp ((branch (vector-ref root 1))
				(shift (- shift *branch-bits*)))
			 (if (= shift *branch-bits*)
			     (writable-branch branch edit)
			     (lp (vector-ref branch 0)
				 (- shift *branch-bits*))))))
	     (set-transient-fector-tail! fector tail)
	     (set-transient-fector-root! fector (vector-ref root 0))
	     (set-transient-fector-shift! fector (- shift *branch-bits*))))
	  (else
	   (let ((root* (writable-branch root edit)))
	     (unless (eq? root root*)
	       (set-transient-fector-root! fector root*))
	     (let ((tail (pop-tail! (- length 1 *branch-size*) shift root*)))
	       (set-transient-fector-tail! fector tail))))))))
     (set-transient-fector-length! fector (1- length))
     fector)
    (($ <fector>)
     (fector-pop! (transient-fector fector)))))

(define (fector-pop fector)
  (define (pop-tail i shift root)
    (let* ((shift (- shift *branch-bits*))
	   (idx (logand (ash i (- shift)) *branch-mask*)))
      (call-with-values (lambda ()
			  (if (= shift *branch-bits*)
			      (values #f (vector-ref root idx))
			      (pop-tail i shift (vector-ref root idx))))
	(lambda (branch tail)
	  (values (clone-branch-and-set root idx branch)
		  tail)))))
  (match fector
    (($ <fector> length shift root tail)
     (let* ((tail-length (compute-tail-length length)))
       (cond
	((< 1 tail-length)
	 ;; Normal case: just trim the tail.
	 (when (zero? length)
	   (error "can't pop from empty fector"))
	 (let ((new-tail (make-vector (1- tail-length) #f)))
	   (vector-move-left! tail 0 (1- tail-length) new-tail 0)
	   (make-fector (1- length) shift root new-tail)))
	(else
	 ;; Tail array will be empty; pop a tail from the tree.
	 (cond
	  ((= length (1+ *branch-size*))
	   ;; The root becomes the tail.
	   (make-fector *branch-size* *branch-bits* *empty-branch* root))
	  ((= length (+ (ash 1 shift) *branch-size* 1))
	   ;; Shrink the tree.
	   (let ((tail (let lp ((branch (vector-ref root 1))
				(shift (- shift *branch-bits*)))
			 (if (= shift *branch-bits*)
			     branch
			     (lp (vector-ref branch 0)
				 (- shift *branch-bits*))))))
	     (make-fector (1- length) (- shift *branch-bits*)
			  (vector-ref root 0) tail)))
	  (else
	   (call-with-values (lambda ()
			       (pop-tail (- length 1 *branch-size*) shift root))
	     (lambda (root tail)
	       (make-fector (1- length) shift root tail)))))))))
    (($ <transient-fector>)
     (fector-pop (persistent-fector fector)))))

(define (fector-ref fector i)
  (define (ref length shift root tail)
    (unless (< i length)
      (error "index out of range" i))
    (let ((tail-offset (compute-tail-offset length)))
      (cond
       ((<= tail-offset i) (vector-ref tail (- i tail-offset)))
       (else
	(let lp ((shift shift) (root root))
	  (let* ((shift (- shift *branch-bits*)))
	    (if (zero? shift)
		(vector-ref root (logand i *branch-mask*))
		(let ((idx (logand (ash i (- shift)) *branch-mask*)))
		  (lp shift (vector-ref root idx))))))))))
  (match fector
    (($ <fector> length shift root tail)
     (ref length shift root tail))
    (($ <transient-fector> length shift root tail edit)
     (assert-readable! edit)
     (ref length shift root tail))))

(define (fector-fold f fector seed)
  (define (visit-branch branch start end shift seed)
    (let* ((shift (- shift *branch-bits*))
	   (end-idx (ash (- end start 1) (- shift)))
	   (inc (ash 1 shift)))
      (let lp ((i 0) (start start) (seed seed))
        (let ((node (vector-ref branch i)))
          (if (< i end-idx)
              (lp (1+ i)
                  (+ start inc)
                  (if (= shift *branch-bits*)
                      (visit-leaf node start seed)
                      (visit-branch node start (+ start inc) shift seed)))
              (if (= shift *branch-bits*)
                  (visit-leaf node start seed)
                  (visit-branch node start end shift seed)))))))
  (define (visit-leaf leaf start seed)
    (let lp ((n 0) (seed seed))
      (if (< n *branch-size*)
	  (lp (1+ n) (f (+ start n) (vector-ref leaf n) seed))
	  seed)))
  (match fector
    (($ <fector> length shift root tail)
     (let* ((tail-offset (pk (compute-tail-offset length)))
	    (seed (if (<= length *branch-size*)
		      seed
		      (if (= shift *branch-bits*)
			  (visit-leaf root 0 seed)
			  (visit-branch root 0 tail-offset shift seed)))))
       (let ((tail-length (- length tail-offset)))
	 (let lp ((n 0) (seed seed))
	   (let ((idx (+ tail-offset n)))
	     (if (= idx length)
		 seed
		 (lp (1+ n) (f idx (vector-ref tail n) seed))))))))
    (($ <transient-fector>)
     (fector-fold f (persistent-fector fector) seed))))

(define (fector-fold-right f fector seed)
  (define (visit-branch branch start end shift seed)
    (let* ((shift (- shift *branch-bits*))
	   (end-idx (ash (- end start 1) (- shift)))
	   (inc (ash 1 shift)))
      (let lp ((i end-idx)
               (end (+ start (* end-idx inc)))
               (seed (let ((node (vector-ref branch end-idx))
                           (start (+ start (* end-idx inc))))
                       (if (= shift *branch-bits*)
                           (visit-leaf node start seed)
                           (visit-branch node start end shift seed)))))
	(if (zero? i)
	    seed
	    (let ((i (1- i)))
	      (let ((node (vector-ref branch i))
		    (start (- end inc)))
		(lp i
		    start
		    (if (= shift *branch-bits*)
			(visit-leaf node start seed)
			(visit-branch node start end shift seed)))))))))
  (define (visit-leaf leaf start seed)
    (let lp ((n *branch-size*) (seed seed))
      (if (zero? n)
	  seed
	  (let ((n (1- n)))
	    (lp n (f (+ start n) (vector-ref leaf n) seed))))))
  (match fector
    (($ <fector> length shift root tail)
     (let* ((tail-offset (compute-tail-offset length)))
       (let lp ((i (1- length)) (seed seed))
	 (if (<= tail-offset i)
	     (lp (1- i) (f i (vector-ref tail (- i tail-offset)) seed))
	     (if (<= length *branch-size*)
		 seed
		 (if (= shift *branch-bits*)
		     (visit-leaf root 0 seed)
		     (visit-branch root 0 tail-offset shift seed)))))))
    (($ <transient-fector>)
     (fector-fold f (persistent-fector fector) seed))))

(define (benchmark-persistent-append count)
  (let lp ((n 0) (fec empty-fector))
    (if (< n count)
	(lp (1+ n) (fector-push fec n))
	fec)))

(define (benchmark-transient-append count)
  (let lp ((n 0) (fec empty-fector))
    (if (< n count)
	(lp (1+ n) (fector-push! fec n))
	(persistent-fector fec))))