Yet another parametric type in Common Lisp

yapt.lisp

(eval-when (:compile-toplevel :load-toplevel :execute)
  (ql:quickload :fiveam)
  (use-package '(:fiveam :sb-mop)))

(defclass parametric-class (class)
  ((parameters :initform nil :initarg :parameters)
   (cache-reified-classes :initform (make-hash-table :test #'equalp)
                          :initarg :cache-reified-classes)))
(defclass parametric-structure-class (structure-class parametric-class)
  ((original-direct-slots :initform nil :initarg :original-direct-slots)))

(defclass parametric-standard-class (standard-class parametric-class) ())

(deftype type-specifier () '(or symbol list class))

(declaim (ftype (function (parametric-class list) (values parametric-class &optional)) make-reified-class))
(defun make-reified-class (class parameters)
  (declare (parametric-class class)
           (list parameters))
  (when (slot-value class 'parameters)
    (error "MAKE-REIFIED-CLASS was called for the already reified class ~A." class))
  (let ((cache (slot-value class 'cache-reified-classes)))
    (or (gethash parameters cache)      
        (let* ((class-name (class-name class))
               (reified-class-name (gensym (symbol-name class-name)))
               (reified-class
                 (ensure-class reified-class-name
                               :metaclass (class-of class)
                               :direct-superclasses `(,class-name)
                               :parameters parameters
                               :cache-reified-classes (slot-value class 'cache-reified-classes))))
          (setf (gethash parameters cache) reified-class)))))

(defun same-typep (type-specifier1 type-specifier2 &optional env)
  (and (subtypep type-specifier1 type-specifier2 env)
       (subtypep type-specifier2 type-specifier1 env)))

(defun extract-parameter (ctype parametric-class template)
  (declare (sb-kernel:ctype ctype)
           (parametric-class parametric-class)
           (symbol template))
  (labels ((classoid-parameter (classoid)
             (let ((class (find-class (sb-kernel:classoid-name classoid))))
               (when (subtypep class parametric-class)
                 (cdr (assoc template (slot-value class 'parameters)))))))
    (typecase ctype
      (sb-kernel:classoid (classoid-parameter ctype))
      (sb-kernel:union-type
       (reduce (lambda (p1 p2) (if (same-typep p1 p2) p1 nil))
               (sb-kernel:union-type-types ctype)
               :key (lambda (ctype) (extract-parameter ctype parametric-class template))))
      (sb-kernel:intersection-type
       (reduce (lambda (p1 p2)
                 (if (and p1 p2)
                     (if (same-typep p1 p2)
                         p1
                         (error "Contradictory parameters derived: ~A ~A. Note that yapt's type parameters are invariant." p1 p2))
                     (or p1 p2)))
               (sb-kernel:intersection-type-types ctype)
               :key (lambda (ctype) (extract-parameter ctype parametric-class template))))
      (t nil))))

;; Example parametric structure (binary heap)
(progn
  (when (find-class 'heap nil)
    (setf (find-class 'heap nil) nil))
  (defstruct (heap (:constructor %make-heap
                       (size
                        &key test (element-type t)
                        &aux (data (make-array (1+ size) :element-type element-type))))
                   (:conc-name %heap-)
                   (:copier nil))
    (data #() :type (simple-array * (*)))
    (test #'< :type function)
    (next-position 1 :type (integer 1 #.most-positive-fixnum)))
  (sb-mop:ensure-class-using-class nil 'heap
                                   :metaclass 'parametric-structure-class
                                   :original-direct-slots (class-direct-slots (find-class 'heap))))

(defmethod print-object ((x heap) stream)
  (sb-kernel:default-structure-print x stream sb-kernel:*current-level-in-print*))

(defmethod make-load-form ((object parametric-class) &optional env)
  (declare (ignore env))
  `(make-reified-class (find-class 'heap) ',(slot-value object 'parameters)))

;; (declaim (inline make-heap))
(sb-c:defknown make-heap ((integer 0 #.array-total-size-limit) &key (:test function) (:element-type type-specifier))
    heap ()
  :overwrite-fndb-silently t)

;; (declaim (ftype (function ((integer 0 #.array-total-size-limit) &key (:test function) (:element-type type-specifier)) (values heap &optional)) make-heap))

(defun make-heap (size &key (test #'<) (element-type t) &aux (data (make-array (1+ size) :element-type element-type)))
  (make-instance (make-reified-class (find-class 'heap) `((element-type . ,element-type)))
                 :data data :test test))

(sb-c:deftransform make-heap ((size &key (test #'<) element-type) (* &key (:test *) (:element-type *)) *)
  (unless (sb-c::constant-lvar-p element-type)
    (sb-c::give-up-ir1-transform "ELEMENT-TYPE is not constant."))
  (let* ((reified-class (make-reified-class (find-class 'heap)
                                            `((element-type . ,(sb-c::lvar-value element-type))))))
    ;; (describe reified-class)
    `(let ((data (make-array (1+ size) :element-type element-type)))
       (the ,reified-class (make-instance ,reified-class :data data :test test)))))

;; (sb-c:defoptimizer (make-heap sb-c:derive-type) ((size &key test element-type))
;;   (declare (ignore size test))
;;   (format t "optimizer called ~A~%" element-type)
;;   (if (sb-c::constant-lvar-p element-type)
;;       (sb-c::specifier-type (make-reified-class (find-class 'heap) `((element-type . ,(sb-c::lvar-value element-type)))))
;;       (sb-c::specifier-type 'heap)))

;; (defmethod initialize-instance ((instance heap) &rest initargs)
;;   (apply #'shared-initialize instance t initargs))

(defmethod initialize-instance ((instance heap) &rest initargs &key (data #()) (test #'<) (next-position 1) &allow-other-keys)
  (declare (ignore initargs))
  (setf (%heap-data instance) data
        (%heap-test instance) test
        (%heap-next-position instance) next-position)
  instance)

(deftype specialized-heap (&optional element-type)
  (if (or (null element-type) (eql '* element-type))
      (find-class 'heap)
      (make-reified-class (find-class 'heap) `((element-type . ,element-type)))))

(sb-c:defknown heap-data (heap) (simple-array * (*))
    (sb-c:foldable sb-c:flushable)
  :overwrite-fndb-silently t)

(defun heap-data (instance)
  (%heap-data instance))

(sb-c:deftransform heap-data ((instance) (heap))
  (let* ((ctype (sb-c::lvar-type instance))
         (element-type (or (extract-parameter ctype (find-class 'heap) 'element-type) '*)))
    `(the (simple-array ,element-type (*)) (%heap-data instance))))

(declaim (inline heap-test))
(defun heap-test (instance)
  (%heap-test instance))

(declaim (inline heap-next-position))
(defun heap-next-position (instance)
  (%heap-next-position instance))

(declaim (inline (setf heap-next-position)))
(defun (setf heap-next-position) (value instance)
  (setf (%heap-next-position instance) value))

(declaim (inline heap-push))
(defun heap-push (obj heap)
  (symbol-macrolet ((next-position (heap-next-position heap)))
    (let ((data (heap-data heap))
          (test (heap-test heap)))
      (labels ((update (pos)
                 (unless (= pos 1)
                   (let ((parent-pos (ash pos -1)))
                     (when (funcall test (aref data pos) (aref data parent-pos))
                       (rotatef (aref data pos) (aref data parent-pos))
                       (update parent-pos))))))
        (setf (aref data next-position) obj)
        (update next-position)
        (incf next-position)
        heap))))

(declaim (inline heap-pop))
(defun heap-pop (heap &optional (error t) null-value)
  (symbol-macrolet ((next-position (heap-next-position heap)))
    (let ((data (heap-data heap))
          (test (heap-test heap)))
      (labels ((update (pos)
                 (declare ((integer 1 #.most-positive-fixnum) pos))
                 (let* ((child-pos1 (+ pos pos))
                        (child-pos2 (1+ child-pos1)))
                   (when (<= child-pos1 next-position)
                     (if (<= child-pos2 next-position)
                         (if (funcall test (aref data child-pos1) (aref data child-pos2))
                             (unless (funcall test (aref data pos) (aref data child-pos1))
                               (rotatef (aref data pos) (aref data child-pos1))
                               (update child-pos1))
                             (unless (funcall test (aref data pos) (aref data child-pos2))
                               (rotatef (aref data pos) (aref data child-pos2))
                               (update child-pos2)))
                         (unless (funcall test (aref data pos) (aref data child-pos1))
                           (rotatef (aref data pos) (aref data child-pos1))))))))
        (if (= next-position 1)
            (if error
                (error "No element in heap")
                null-value)
            (prog1 (aref data 1)
              (decf next-position)
              (setf (aref data 1) (aref data next-position))
              (update 1)))))))

(declaim (inline heap-peak))
(defun heap-peak (heap &optional (error t) null-value)
  (if (= 1 (heap-next-position heap))
      (if error
          (error "No element in heap")
          null-value)
      (aref (heap-data heap) 1)))

(defun bench (&optional (size 2000000))
  (declare (optimize (speed 3)))
  (let* ((heap (make-heap size :element-type 'fixnum))
         (seed (seed-random-state 0)))
    (declare ((specialized-heap fixnum) heap))
    (time (dotimes (i size)
            (heap-push (random most-positive-fixnum seed) heap)))
    (time (dotimes (i size)
            (heap-pop heap)))))

(defun bench-init (&optional (num 10000000))
  (declare (optimize (speed 3))
           (fixnum num))
  (time (dotimes (i num) (make-heap 1 :element-type 'fixnum))))

;; Evaluation took:
;;   3.625 seconds of real time
;;   3.609375 seconds of total run time (3.609375 user, 0.000000 system)
;;   [ Run times consist of 0.032 seconds GC time, and 3.578 seconds non-GC time. ]
;;   99.56% CPU
;;   3 lambdas converted
;;   9,058,020,173 processor cycles
;;   1,760,201,280 bytes consed

;; Evaluation took:
;;   3.687 seconds of real time
;;   3.687500 seconds of total run time (3.687500 user, 0.000000 system)
;;   [ Run times consist of 0.062 seconds GC time, and 3.626 seconds non-GC time. ]
;;   100.03% CPU
;;   9,193,560,006 processor cycles
;;   1,599,996,000 bytes consed
  
;; For test

(test heap-test
  (let ((h (make-heap 20)))
    (finishes (dolist (o (list 7 18 22 15 27 9 11))
                (heap-push o h)))
    (is (= 7 (heap-peak h)))
    (is (equal '(7 9 11 15 18 22 27)
               (loop repeat 7 collect (heap-pop h))))
    (signals error (heap-pop h))
    (is (eql 'eof (heap-pop h nil 'eof)))
    (is (eql 'eof (heap-peak h nil 'eof))))
  (is (typep (heap-data (make-heap 10 :element-type 'fixnum))
             '(simple-array fixnum (*)))))

(run! 'heap-test)

(test parameter-test
  (is (eql 'long-float
           (extract-parameter (sb-c::specifier-type '(and (specialized-heap long-float) (satisfies heap-p)))
                            (find-class 'heap)
                            'element-type))))
(run! 'parameter-test)

;; Memo:
;; - Cached reified classs should be looked up with SAME-TYPEP.