a leftist heap implementation and incremental wheel sieves built on top of it

leftist-heap-wheel-sieves.clj

(declare leftist-heap)

(defprotocol heap
  (heap-empty? [h])
  (heap-merge [h1 h2])
  (heap-rank [h])
  (heap-insert [h x])
  (heap-get-min [h])
  (heap-del-min [h]))

(deftype empty-leftist-heap [<=?] :as self
  heap
  (heap-empty? [] true)
  (heap-merge [h] h)
  (heap-rank [] 0)
  (heap-insert [x] (leftist-heap 1 x self self <=?))
  (heap-get-min [] (throw (Exception. "getting minimum of empty leftist heap")))
  (heap-del-min [] (throw (Exception. "deleting minimum of empty leftist heap"))))

(deftype leftist-heap [rank elem left right <=?] :as self
  heap
  (heap-empty? [] false)
  (heap-merge [h]
    (if (heap-empty? h)
      self
      (let [other (heap-get-min h)
            make (fn [x ha hb]
                   (if (< (heap-rank ha) (heap-rank hb))
                     (leftist-heap. (inc (heap-rank ha))
                                    x
                                    hb
                                    ha
                                    <=?)
                     (leftist-heap. (inc (heap-rank hb))
                                    x
                                    ha
                                    hb
                                    <=?)))]
        (if (<=? elem other)
          (make elem left (heap-merge right h))
          (make other
                (:left h)
                (heap-merge self
                            (:right h)))))))
  (heap-rank [] rank)
  (heap-insert [x]
    (heap-merge
     (leftist-heap. 1 x (empty-leftist-heap <=?) (empty-leftist-heap <=?) <=?)
     self))
  (heap-get-min [] elem)
  (heap-del-min [] (heap-merge left right)))

(defprotocol heap-map
  (heap-map-empty? [h])
  (heap-map-merge [h1 h2])
  (heap-map-rank [h])
  (heap-map-insert [h k v])
  (heap-map-get-min-key [h])
  (heap-map-get-min-val [h])
  (heap-map-del-min [h]))

(def leftist-heap-map)

(deftype empty-leftist-heap-map [] :as self
  heap-map
  (heap-map-empty? [] true)
  (heap-map-merge [h] h)
  (heap-map-rank [] 0)
  (heap-map-insert [k v] (leftist-heap-map 1 k v self self))
  (heap-map-get-min-key [] (throw (Exception. "getting minimum of empty leftist heap-map")))
  (heap-map-get-min-val [] (throw (Exception. "getting minimum of empty leftist heap-map")))
  (heap-map-del-min [] (throw (Exception. "deleting minimum of empty leftist heap-map"))))

(deftype leftist-heap-map [rank k v left right] :as self
  heap-map
  (heap-map-empty? [] false)
  (heap-map-merge [h]
    (if (heap-map-empty? h)
      self
      (let [other (:k h)
            make (fn [k v ha hb]
                   (if (< (heap-map-rank ha) (heap-map-rank hb))
                     (leftist-heap-map. (inc (heap-map-rank ha))
                                        k
                                        v
                                        hb
                                        ha)
                     (leftist-heap-map. (inc (heap-map-rank hb))
                                        k
                                        v
                                        ha
                                        hb)))]
        (if (< k other)
          (make k v left (heap-map-merge right h))
          (make other
                (:v h)
                (:left h)
                (heap-map-merge self
                            (:right h)))))))
  (heap-map-rank [] rank)
  (heap-map-insert [k v]
    (heap-map-merge
     (leftist-heap-map. 1 k v (empty-leftist-heap-map) (empty-leftist-heap-map))
     self))
  (heap-map-get-min-key [] k)
  (heap-map-get-min-val [] v)
  (heap-map-del-min [] (heap-map-merge left right)))

(defn make-wheel [n]
  (let [ps (take (inc n) (incremental-sieve))
        init (butlast ps)
        start (last ps)
        m (reduce * init)
        not-skippable (remove (fn [n] (some #(zero? (rem n %)) init))
                              (range start (inc (+ start m))))
        steps (first (reduce (fn [[steps previous] current]
                               [(conj steps (- current previous)) current])
                             [[] (first not-skippable)]
                             (rest not-skippable)))]
    [ps steps]))

(defn lheap-incremental-wheel-sieve [wheel & wheel-args]
  (let [<=? (fn [[x _] [y _]] (<= x y))
        [ps steps] (if (fn? wheel) (apply wheel wheel-args) wheel)
        start (last ps)
        steps (cycle steps)
        next (fn next [crossouts candidate steps]
               (let [[c p] (heap-get-min crossouts)]
                 (cond
                   (== c candidate) (next (-> crossouts
                                              heap-del-min
                                              (heap-insert [(+ c p) p]))
                                          (+ candidate (first steps))
                                          (rest steps))
                   (< c candidate) (next (-> crossouts
                                             heap-del-min
                                             (heap-insert [(+ c p) p]))
                                         candidate
                                         steps)
                   :else (cons candidate
                               (lazy-seq
                                 (next
                                  (heap-insert crossouts
                                               [(* candidate candidate) candidate])
                                  (+ candidate (first steps))
                                  (rest steps)))))))]
    (concat (butlast ps)
            (lazy-seq (next (reduce (fn [h p]
                                      (heap-insert h [(* p p) p]))
                                    (empty-leftist-heap <=?)
                                    (butlast ps))
                            start
                            steps)))))

(defn lheap-map-incremental-wheel-sieve [wheel & wheel-args]
  (let [[ps steps] (if (fn? wheel) (apply wheel wheel-args) wheel)
        start (last ps)
        steps (cycle steps)
        next (fn next [crossouts candidate steps]
               (let [c (heap-map-get-min-key crossouts)
                     p (heap-map-get-min-val crossouts)]
                 (cond
                   (== c candidate) (next (-> crossouts
                                              heap-map-del-min
                                              (heap-map-insert (+ c p) p))
                                          (+ candidate (first steps))
                                          (rest steps))
                   (< c candidate) (next (-> crossouts
                                             heap-map-del-min
                                             (heap-map-insert (+ c p) p))
                                         candidate
                                         steps)
                   :else (cons candidate
                               (lazy-seq
                                 (next
                                  (heap-map-insert crossouts
                                                   (* candidate candidate)
                                                   candidate)
                                  (+ candidate (first steps))
                                  (rest steps)))))))]
    (concat (butlast ps)
            (lazy-seq (next (reduce (fn [h p]
                                      (heap-map-insert h (* p p) p))
                                    (empty-leftist-heap-map)
                                    (butlast ps))
                            start
                            steps)))))