; SPDX-License-Identifier: LGPL-3.0-or-later
; Copyright (C) 2010-2021 Tony Garnock-Jones <tonygarnockjones@gmail.com>

#lang racket/base
;; Bags and Deltas (which are Bags where item-counts can be negative).

(provide make-bag ;; mutable
         bag ;; immutable
         bag-change!
         bag-change
         bag-ref
         bag-clear!
         bag-member?
         bag-empty?
         bag-key-count
         in-bag
         in-bag/count
         for/bag/count
         for/bag
         set->bag
         bag->set)

(require racket/set)

;; A `(MutableBagof X)` is a `(MutableHash X Nat)`, where the `Nat`
;; against an `X` is its replication count in the bag.
;;
;; A `(Bagof X)` is similar, but immutable.
;;
;; `MutableDeltaof` and `Deltaof` are like `MutableBagof` and `Bagof`,
;; respectively, except the replication counts can be negative.

(define make-bag make-hash)
(define bag hash)

(define (bag-change! b x delta)
  (define old-count (bag-ref b x))
  (define new-count (+ old-count delta))
  (if (zero? new-count)
      (begin (hash-remove! b x)
             (if (zero? old-count) 'absent->absent 'present->absent))
      (begin (hash-set! b x new-count)
             (if (zero? old-count) 'absent->present 'present->present))))

(define (bag-change b x delta #:clamp? [clamp? #f])
  (define old-count (bag-ref b x))
  (define new-count (if clamp?
                        (max 0 (+ old-count delta))
                        (+ old-count delta)))
  (if (zero? new-count)
      (values (hash-remove b x)
              (if (zero? old-count) 'absent->absent 'present->absent))
      (values (hash-set b x new-count)
              (if (zero? old-count) 'absent->present 'present->present))))

(define (bag-ref b x)
  (hash-ref b x 0))

(define bag-clear! hash-clear!)
(define bag-member? hash-has-key?)
(define bag-empty? hash-empty?)
(define bag-key-count hash-count)

(define-syntax-rule (in-bag piece ...) (in-hash-keys piece ...))
(define-syntax-rule (in-bag/count piece ...) (in-hash piece ...))

(define-syntax-rule (for/bag/count iters expr ...) (for/hash iters expr ...))
(define-syntax-rule (for/bag iters expr ...) (for/bag/count iters (values (begin expr ...) 1)))

(define (set->bag s [count 1])
  (for/hash [(e (in-set s))]
    (values e count)))

(define (bag->set b)
  (list->set (hash-keys b)))