reservoir sampling

reservoir-sample.lisp

;; Implements reservoir sampling
;; https://www.omniref.com/ruby/gems/stream_sampler/0.0.1/symbols/StreamSampler?#annotation=4094626&line=2

(defpackage #:snippets/reservoir-sample
  (:use #:cl)
  (:export
   #:no-more
   #:reservoir-sample))

(in-package #:snippets/reservoir-sample)

(defvar *iterable-terminator* (cons :terminator nil))

(define-symbol-macro no-more *iterable-terminator*)

(defun iterable-function (iterable)
  (etypecase iterable
    (sequence (seq-iterable iterable))
    (symbol (fdefinition iterable))
    (function iterable)))

(defun map-iterable (function iterable)
  (setf iterable (iterable-function iterable))
  (do ((index 0 (1+ index))
       (element (funcall iterable)
                (funcall iterable)))
      ((eq element no-more))
    (funcall function element index)))

(defmacro do-iterable (((element &optional index) iterable) &body forms)
  (when (null index) (setf index (gensym)))
  `(block nil
     (map-iterable (lambda (,element ,index)
                     (declare (ignorable ,element ,index))
                     ,@forms)
                   ,iterable)))

(defun list-iterable (list)
  (lambda ()
    (cond ((null list) no-more)
          (t (pop list)))))

(defun vector-iterable (vector)
  (let ((i 0))
    (lambda ()
      (cond ((>= i (length vector)) no-more)
            (t (prog1 (aref vector i)
                 (incf i)))))))

(defun seq-iterable (sequence)
  (if (vectorp sequence)
      (vector-iterable sequence)
      (list-iterable sequence)))

(defun reservoir-sample (iterable n)
  (let ((sample (make-array n :initial-element nil)))
    (when (plusp n)
      (do-iterable ((element index) iterable)
        (if (< index n)
            (setf (aref sample index) element)
            (let ((j (random index)))
              (if (< j n)
                  (setf (aref sample j) element))))))
    sample))