Lazy views over n merged Bifurcan collections, assuming disjoint keys for sets/maps.

bifurcan-lazy-union.clj

(ns bifurcan-lazy-union
  (:require [bifurcan-clj [core :as b]
                          [list :as bl]
                          [map :as bm]
                          [set :as bs]]
            [dom-top.core :refer [loopr]]
            [potemkin :refer [definterface+]])
  (:import (io.lacuna.bifurcan ISet
                               Lists
                               Sets
                               Maps)
           (java.util OptionalLong)
           (java.util.function Function
                               LongFunction)))


(defn lazy-list-concat
  "Lazy list concat for Bifurcan."
  [lists]
  (case (count lists)
    0 bl/empty
    1 (first lists)
    (let [size (reduce + (map b/size lists))
          nth  (reify LongFunction
                 (apply [_ i]
                   (loopr [i i]
                          [list lists]
                          (let [size (b/size list)]
                            (if (< i size)
                              ; It's in this list
                              (b/nth list i)
                              ; Try the next one
                              (recur (- i size)))))))]
          (Lists/from size nth))))

(defn lazy-set-union
  "Lazy set union for Bifurcan. ONLY works if the sets are disjoint--for our
  purposes, they always are."
  [sets]
  (case (count sets)
    0 bs/empty
    1 (first sets)
    (let [; List of elements
          size (reduce + (map b/size sets))
          nth  (reify LongFunction
                 (apply [_ i]
                   (loopr [i i]
                          [set sets]
                          (let [size (b/size set)]
                            (if (< i size)
                              ; It's in this set
                              (b/nth set i)
                              ; Try the next one
                              (recur (- i size)))))))
          elements (Lists/from size nth)
          ; Function for the index of an element
          index-of (reify Function
                     (apply [_ x]
                       (loopr [i 0] ; Index of first element in set
                              [set sets]
                              (if-let [j (bs/index-of set x)]
                                (OptionalLong/of (+ i j))
                                (recur (+ i (b/size set))))
                              (OptionalLong/empty))))]
      (Sets/from elements, index-of))))

(defn lazy-map-union
  "Lazy map union for Bifurcan. ONLY works if the maps are disjoint--for our
  purposes, they always are."
  [maps]
  (let [ks (lazy-set-union (mapv bm/keys maps))
        to-v (reify Function
               (apply [_ k]
                 (loopr []
                        [m maps]
                        (let [v (bm/get m k ::not-found)]
                          (if (identical? ::not-found v)
                            (recur)
                            v))
                        (throw (IllegalStateException.
                                 (str "Expected a value for key " k))))))]
    (Maps/from ks to-v)))