path: root/mcgoron/weight-balanced-trees/srfi/113/sets.scm

#| Copyright 2024 Peter McGoron
 |
 | Licensed under the Apache License, Version 2.0 (the "License");
 | you may not use this file except in compliance with the License.
 | You may obtain a copy of the License at
 |
 |     http://www.apache.org/licenses/LICENSE-2.0
 |
 | Unless required by applicable law or agreed to in writing, software
 | distributed under the License is distributed on an "AS IS" BASIS,
 | WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 | See the License for the specific language governing permissions and
 | limitations under the License.
 |#

(define-record-type <set>
  (raw-set comparator node)
  set?
  (comparator set-element-comparator)
  (node get-node))

(define (check-compatible set1 set2)
  (let ((val (binary-compatible set1 set2)))
    (if (not val)
        (error "sets have different comparators" set1 set2)
        val)))

;;; ;;;;;;;;;;;;;;;;
;;; Constructors
;;; ;;;;;;;;;;;;;;;;

(define (set-unfold comparator stop? mapper successor seed)
  (let loop ((set '())
             (seed seed))
    (if (stop? seed)
        (raw-set comparator set)
        (let ((new-value (mapper seed)))
          (loop (insert comparator set new-value)
                (successor seed))))))

(define (set comparator . elements)
  (list->set comparator elements))

;;; ;;;;;;;;;;;;;;;;;
;;; Predicates (besides set?)
;;; ;;;;;;;;;;;;;;;;;

(define sentinel-value (cons #f #f))

(define (set-contains? set element)
  (not (eq? (search (set-element-comparator set)
                      element
                      (get-node set)
                      (lambda () sentinel-value))
            sentinel-value)))

(define (set-empty? set)
  (null? (get-node set)))

(define (set-disjoint? set1 set2)
  #;(set-empty? (set-intersection set1 set2))
  ;; More optimized version.
  ;; 
  ;; List the values of the sets in order. If any set is exhausted, then
  ;; the sets are disjoint. If any element is equal, then the sets are
  ;; not disjoint.
  ;; 
  ;; If the element from set 1 is less than the element from set 2, then
  ;; get the next element from set 1 (if any) and repeat. Since the
  ;; elements are obtained in order, any elements after the current
  ;; element of set 2 must be greater than the seen elements from set 1.
  (let ((gen1 (set->in-order-generator set1))
        (gen2 (set->in-order-generator set2))
        (cmp (check-compatible set1 set2)))
    (let loop ((value1 (gen1))
               (value2 (gen2)))
      (if (or (eof-object? value1) (eof-object? value2))
          #t
          (comparator-if<=> cmp value1 value2
            (loop (gen1) value2)
            #f
            (loop value1 (gen2)))))))

;;; ;;;;;;;;;;;;;;;;;;;
;;; Accessors
;;; ;;;;;;;;;;;;;;;;;;;

(define (set-member set element default)
  (search (set-element-comparator set)
          element
          (get-node set)
          (lambda () default)))

;;; ;;;;;;;;;;;;;;;;;;;
;;; Updaters
;;; ;;;;;;;;;;;;;;;;;;;

(define (set-adjoin set . elements)
  (set-adjoin-all set elements))
(define set-adjoin! set-adjoin)

(define (set-replace set . elements)
  (set-replace-all elements))
(define set-replace! set-replace)

(define (set-delete-all set elements)
  (let ((cmp (set-element-comparator set)))
    (raw-set
     cmp
     (fold (lambda (element node)
             (delete cmp node element))
           (get-node set)
           elements))))
(define set-delete-all! set-delete-all)

(define (set-delete set . elements)
  (set-delete-all set elements))
(define set-delete! set-delete)

(define (set-search! set element failure success)
  ;; The SRFI mandates that `failure` and `success` are tail-called.
  (define (%insert obj)
    (values (set-adjoin set element) obj))
  (define (%ignore obj)
    (values set obj))
  (define (%remove obj)
    (values (set-remove set element) obj))
  (define (%update new-element obj)
    (values (set-replace (set-remove set element) new-element)
            obj))
  (let ((value (set-member set element (eof-object))))
    (if (eof-object? value)
        (failure %insert %ignore)
        (success value %update %remove))))

(define (set-size set) (get-size (get-node set)))

;;; ;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;
;;; The whole set
;;; ;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;

(define set-find
  (case-lambda
    ((predicate set default-value)
     (set-find predicate set default-value (lambda (x) x)))
    ((predicate set default-value transform)
     (let loop ((queue (list (get-node set))))
       (cond
         ((null? queue) (default-value))
         ((null? (car queue)) (loop (cdr queue)))
         (else (with-node ((car queue) data ("<" left) (">" right))
                 (if (predicate data)
                     (transform data)
                     (loop (cons* left right (cdr queue)))))))))))

(define (set-count predicate set)
  (define (count node)
    (if (null? node)
        0
        (+ (if (predicate (get-data node)) 1 0)
           (count (get-left node))
           (count (get-right node)))))
  (count (get-node set)))

(define (set-any? predicate set)
  (set-find predicate set (lambda () #f) (lambda (x) #t)))

(define (set-every? predicate set)
  (set-find (lambda (x) (not (predicate x)))
            set
            (lambda () #t)
            (lambda (x) #f)))

;;; ;;;;;;;;;;;;;;;;;;;;;;;;
;;; Mapping and folding
;;; ;;;;;;;;;;;;;;;;;;;;;;;;

(define (set-map comparator proc old-set)
  (let ((cmp (set-element-comparator old-set)))
    (raw-set cmp
             (set-fold (lambda (value new-node)
                         (insert cmp new-node (proc value)))
                       '()
                       old-set))))

(define (set-for-each proc set)
  (let loop ((node (get-node set)))
    (when (not (null? node))
      (proc (get-data node))
      (loop (get-left node))
      (loop (get-right node)))))

(define (set-fold proc nil set)
  (define (node-fold nil node)
    (if (null? node)
        nil
        (with-node (node data ("<" left) (">" right))
          (let ((nil (proc data nil)))
            (node-fold (node-fold nil left) right)))))
  (node-fold nil (get-node set)))

(define (set-filter predicate? set)
  (define (loop node)
    (if (null? node)
        '()
        (with-node (node data ("<" left) (">" right))
          (if (predicate? data)
              (join data (loop left) (loop right))
              (join2 (loop left) (loop right))))))
  (raw-set (set-element-comparator set)
           (loop (get-node set))))
(define set-filter! set-filter)

(define (set-remove predicate? set)
  (set-filter (lambda (x)
                (not (predicate? x)))
              set))
(define set-remove! set-remove)

(define (set-partition predicate? set)
  (define (loop node)
    (if (null? node)
        (values '() '())
        (with-node (node data ("<" left) (">" right))
          (let-values (((yes-left no-left)
                        (loop left))
                       ((yes-right no-right)
                        (loop right)))
            (if (predicate? data)
                (values (join data yes-left yes-right)
                        (join2 no-left no-right))
                (values (join2 yes-left yes-right)
                        (join data no-left no-right)))))))
  (let-values (((yes no) (loop (get-node set)))
               ((cmp) (set-element-comparator set)))
    (values (raw-set cmp yes)
            (raw-set cmp no))))

(define set-partition! set-partition)

;;; ;;;;;;;;;;;;;;;;;;;;;;
;;; Copying and conversion
;;; ;;;;;;;;;;;;;;;;;;;;;;

(define (set-copy set)
  ;; NOTE: This function is useless for this implementation because nodes
  ;; can never be modified.
  set
  #;(define (node-copy node)
      (if (null? node)
          '()
          (with-node (node data ("<" left) (">" right))
            (wb-tree-node data (node-copy left) (node-copy right)))))
  #;(raw-set (get-element-comparator set) (node-copy node)))

(define (list->set comparator lst)
  (let loop ((node '())
             (lst lst))
    (if (null? lst)
        (raw-set comparator node)
        (loop (insert comparator node (car lst)) (cdr lst)))))

(define (set->list set)
  (set-fold cons '() set))

(define (list->set! set elements)
  (set-union set (list->set (set-element-comparator set) elements)))

;;; ;;;;;;;;;;;;;;;;;;;
;;; Subsets
;;; ;;;;;;;;;;;;;;;;;;;

(define (apply-nary-predicate binary)
  (lambda (first . rest)
    (let loop ((arg1 first)
               (arg-rest rest))
      (if (null? arg-rest)
          #t
          (let ((arg2 (car arg-rest)))
            (if (binary (check-compatible arg1 arg2) arg1 arg2)
                (loop arg2 (cdr arg-rest))
                #f))))))
(define set=?
  (apply-nary-predicate
   (lambda (cmp set1 set2)
     (or (eq? set1 set2)
         (and (= (set-size set1) (set-size set2))
              (let ((gen1 (set->in-order-generator set1))
                    (gen2 (set->in-order-generator set2)))
                (let loop ((value1 (gen1))
                           (value2 (gen2)))
                  (cond
                    ((and (eof-object? value1) (eof-object? value2)) #t)
                    ((=? cmp value1 value2) (loop (gen1) (gen2)))
                    (else #f)))))))))

(define (binary-set<=? cmp set1 set2)
  (or (eq? set1 set2)
      (and (<= (set-size set1) (set-size set2))
           (set-every? (cut set-contains? set2 <>) set1))))

(define set<=?
  (apply-nary-predicate binary-set<=?))

(define (binary-set<? cmp set1 set2)
  (and (not (eq? set1 set2))
       (< (set-size set1) (set-size set2))
       (set-every? (cut set-contains? set2 <>) set1)))

(define set<?
  (apply-nary-predicate binary-set<?))

(define set>?
  (apply-nary-predicate
   (lambda (cmp set1 set2) (binary-set<? cmp set2 set1))))

(define set>=?
  (apply-nary-predicate
   (lambda (cmp set1 set2) (binary-set<=? cmp set2 set1))))

;;; ;;;;;;;;;;;;;;;;
;;; Set theory operations
;;; ;;;;;;;;;;;;;;;;

(define (apply-nary-procedure binary)
  (lambda (first . rest)
    (let ((cmp (set-element-comparator first)))
      (let loop ((arg1 first)
                 (arg-rest rest))
        (if (null? arg-rest)
            arg1
            (let ((arg2 (car arg-rest)))
              (check-compatible arg1 arg2)
              (loop (binary cmp arg1 arg2)
                    (cdr arg-rest))))))))

(define (convert-binary-procedure proc)
  (apply-nary-procedure
   (lambda (cmp arg1 arg2)
     (raw-set cmp (proc cmp (get-node arg1) (get-node arg2))))))

(define set-union (convert-binary-procedure union))
(define set-union! set-union)

(define set-intersection (convert-binary-procedure
                          (lambda (cmp node1 node2)
                            (if (eq? node1 node2)
                                node1
                                (intersection cmp node1 node2)))))
(define set-intersection! set-intersection)

(define set-difference (convert-binary-procedure difference))
(define set-difference! set-difference)

(define set-xor (apply-nary-procedure xor))
(define set-xor! set-xor)

;;; ;;;;;;;;;;;;
;;; exported extensions
;;; ;;;;;;;;;;;;

(define (set-adjoin-all set elements)
  (let ((cmp (set-element-comparator set)))
    (raw-set
     cmp
     (fold (lambda (new set)
             (update cmp
                     set
                     new
                     (lambda (old) old)
                     (lambda () (wb-tree-node new '() '()))))
           (get-node set)
           elements))))

(define (set-replace-all set elements)
  (let ((cmp (set-element-comparator set)))
    (fold (lambda (new set)
            (update cmp
                    set
                    new
                    (lambda (old) new)
                    (lambda ()
                      (wb-tree-node new '() '()))))
          (get-node set)
          elements)))

(define (generator->set comparator gen)
  (raw-set comparator (generator->node comparator gen)))

(define (set->generator set)
  (node->generator (get-node set)))

(define (set->in-order-generator set)
  (node->in-order-generator (get-node set)))

(define (set->reverse-order-generator set)
  (node->reverse-order-generator (get-node set)))

(define (in-order->set comparator container ref length)
  (raw-set comparator
           (in-order-container->node container ref length)))

(define (binary-compatible s1 s2)
  (let ((cmp (set-element-comparator s1)))
    (and (eq? cmp (set-element-comparator s2))
         cmp)))

(define compatible-sets?
  (apply-nary-predicate (lambda (cut set1 set2)
                          (binary-compatible set1 set2))))