;;; -*- Mode: Lisp; Package: FSet; Syntax: ANSI-Common-Lisp -*-

(in-package :fset)

;;; File: fset.lisp
;;; Contents: Top level of FSet, the fast functional set-theoretic datatypes package.
;;;
;;; This file is part of FSet.  Copyright (c) 2007 Sympoiesis, Inc.
;;; FSet is licensed under the Lisp Lesser GNU Public License, or LLGPL.
;;; See: http://opensource.franz.com/preamble.html
;;; This license provides NO WARRANTY.


;;; ================================================================================
;;; Generic functions

;;; We make almost all the interface operations generic to support the addition of
;;; new implementations.

(defgeneric empty? (collection)
  (:documentation "Returns true iff the collection is empty."))

(defgeneric size (collection)
  (:documentation
    "Returns the number of members in a set, seq, or bag, or the number of
pairs in a map.  The size of a bag is the sum of the multiplicities."))

(defgeneric set-size (bag)
  (:documentation
    "Returns the number of unique members in the bag."))

(defgeneric arb (collection)
  (:documentation
    "Returns an arbitrary member or pair of a set, bag, or map.  Specifically,
on a nonempty set, returns two values, an arbitrary member of the set and
true; on a nonempty bag, returns an arbitrary member, its multiplicity,
and true; on a nonempty map, returns an arbitrary domain member, its
associated value, and true.  On an empty set, bag, or map, returns false for
all values.   Please note that \"arbitrary\" does not mean \"randomly selected\";
it simply means that the sole postcondition is that the returned value or pair
is a member of the collection."))

(defgeneric member? (x collection)
  (:documentation
    "Returns true iff `x' is a member of a set, bag, or seq, or (for convenience)
a member of the domain of a map.  Note that for a seq, a linear search is
required."))

(defgeneric multiplicity (x bag)
  (:documentation "Returns the multiplicity of `x' in the bag."))

(defgeneric least (collection)
  (:documentation
    "On a set, returns two values: the smallest member of the set and
true; on a bag, returns three values: the smallest member of the bag, its
multiplicity, and true; on a map, also returns three values: the smallest key
of the map, its value, and true.  If there is not a unique smallest member,
chooses one of the smallest members arbitrarily.  Returns `nil' if the
collection is empty."))

(defgeneric greatest (collection)
  (:documentation
    "On a set, returns two values: the greatest member of the set and
true; on a bag, returns three values: the greatest member of the bag, its
multiplicity, and true; on a map, also returns three values: the greatest key
of the map, its value, and true.  If there is not a unique greatest member,
chooses one of the greatest members arbitrarily.  Returns `nil' if the
collection is empty."))

(defgeneric lookup (collection key)
  (:documentation
    "If `collection' is a map, returns the value to which `key' is mapped.
If `collection' is a seq, takes `key' as an index and returns the
corresponding member (0-origin, of course).  If `collection' is a set or
bag that contains a member `equal?' to `key', returns true and the member
as two values, else false and `nil'; this is useful for canonicalization."))

(defmacro @ (fn-or-collection &rest args)
  "A little hack with two purposes: (1) to make it easy to make FSet maps
behave like Lisp functions in certain contexts; and (2) to somewhat lessen the
pain of writing higher-order code in a two-namespace Lisp like Common Lisp.
The idea is that you can write `(@ fn arg)', and if `fn' is a Lisp function,
it will be funcalled on the argument; otherwise `lookup' (q.v.) will be called
on `fn' and `arg'.  To allow for `@' to be used in more contexts, it actually
can take any number of `args', though `lookup' always takes exactly one.  Thus
you can write `(@ fn arg1 arg2 ...)' when you just want a shorter name for
`funcall'.  As a matter of style, it is suggested that `@' be used only for
side-effect-free functions.  Also, though this doc string has spoken only of
FSet maps, `@' can be used with any type that `lookup' works on.  Can be used
with `setf', but only on collections, not functions, of course."
  (if (> (length args) 1)
      ;; Hmm.  We _could_ listify `args' and use that as the map key.
      `(funcall ,fn-or-collection . ,args)
    (let ((fn-var (gensym "FN-")))
      `(let ((,fn-var ,fn-or-collection))
	 (if (functionp ,fn-var)
	     (funcall ,fn-var . ,args)
	   ;; We do it this way rather than just `(lookup fn-or-collection (car args))'
	   ;; so that we get the right error when `args' is not of length 1.  If this
	   ;; doesn't get compiled well everyplace we care about, we could test the
	   ;; length and issue the error ourselves (if that helps).
	   (lookup ,fn-var . ,args))))))

(defmacro with (collection val1 &optional (val2 nil val2?))
  "A syntactic convenience.  Expands to a call to `with1' if called with two
arguments, or to `with2' if called with three."
  (if val2? `(with2 ,collection ,val1 ,val2)
    `(with1 ,collection ,val1)))

(defgeneric with1 (collection value)
  (:documentation
    "Adds `value' to a set or bag, returning the updated collection."))

(defgeneric with2 (collection key value)
  (:documentation
    "Adds a mapping from `key' to `value' to a map or seq, returning the
updated collection.  In the seq case, `key' must be in the interval
[0, size(collection)]."))

(defmacro less (collection val1 &optional (val2 nil val2?))
  "A syntactic convenience.  Expands to a call to `less1' if called with two
arguments, or to `less2' if called with three."
  (if val2? `(less2 ,collection ,val1 ,val2)
    `(less1 ,collection ,val1)))

(defgeneric less1 (collection value)
  (:documentation
    "Removes `value' from a set, or the pair whose key is `value' from a
map, or one occurrence of `value' from a bag, or the element whose index
is `value' from a seq (shifting subsequent elements down); returns the
updated collection."))

(defgeneric less2 (collection value count)
  (:documentation
    "Removes `count' occurrences of `value' from a bag, returning the updated
collection."))

(defgeneric union (set-or-bag1 set-or-bag2)
  (:documentation
    "Returns the union of the two sets/bags.  The result is a set if both
arguments are sets; otherwise a bag.  The union of two bags is a bag whose
multiplicity, for any value, is the maximum of its multiplicities in the
two argument bags."))

(defgeneric bag-sum (bag1 bag2)
  (:documentation
    "Returns a bag whose multiplicity, for any value, is the sum of its
multiplicities in the two argument bags."))

(defgeneric intersection (set-or-bag1 set-or-bag2)
  (:documentation
    "Returns the intersection of the two sets/bags.  The result is a bag
if both arguments are bags; otherwise a set.  The intersection of two bags
is the bag whose multiplicity, for any value, is the minimum of its
multiplicities in the two argument bags."))

(defgeneric bag-product (bag1 bag2)
  (:documentation
    "Returns a bag whose multiplicity, for any value, is the product of
its multiplicities in the two argument bags."))

(defgeneric set-difference (set1 set2)
  (:documentation
    "Returns the set difference of set1 and set2, i.e., the set containing
every member of `set1' that is not in `set2'."))

(defgeneric set-difference-2 (set1 set2)
  (:documentation
    "Returns `set1 - set2' and `set2 - set1' as two values."))

(defgeneric bag-difference (bag1 bag2)
  (:documentation
    "Returns a bag whose multiplicity, for any value, is its multiplicity
in `bag1' less that in `bag2', but of course not less than zero."))

(defgeneric subset? (set1 set2)
  (:documentation "Returns true iff `set1' is a subset of `set2'."))

(defgeneric subbag? (bag1 bag2)
  (:documentation
    "Returns true iff `bag1' is a subbag of `bag2', that is, for every
member of `bag1', `bag2' contains the same value with at least the same
multiplicity."))

(defgeneric filter (fn collection)
  (:documentation
    "Returns a new collection containing those members or pairs of `collection'
for which `fn' returns true.  If `collection' is a set, bag, or seq, `fn' is
called with one argument; if a map, `fn' is called with two arguments, the key
and the value (the map-default of the result is that of `collection').  As well
as a Lisp function, `fn' can be a map, or a set (which is treated as mapping
its members to true and everything else to false)."))

(defgeneric image (fn collection)
  (:documentation
    "Returns a new collection containing the result of applying `fn' to each
member of `collection', which may be a set, bag, or seq.  In the bag case,
the multiplicity of each member of the result is the sum of the multiplicities
of the values that `fn' maps to it.  As well as a Lisp function, `fn' can be a
map, or a set (which is treated as mapping its members to true and everything
else to false).  `collection' can also be a map, in which case `fn' must be a
Lisp function of two arguments that returns two values (the map-default of the
result is that of `collection')."))

(defgeneric fold (fn collection &optional initial-value)
  (:documentation
    "Iterates over `collection', maintaining a state S; on each iteration, `fn'
is called on S and the next member of `collection', and the result is used as
the new value of S; finally, returns S.  The first iteration is special: if
`initial-value' is supplied, it is used as the initial S; otherwise, the first
member of `collection' is used as the initial S, and `fn' is not called on this
iteration."))

(defgeneric domain (map)
  (:documentation
    "Returns the domain of the map, that is, the set of keys mapped by the map."))

(defgeneric range (map)
  (:documentation
    "Returns the range of the map, that is, the set of all values to which keys
are mapped by the map."))

(defgeneric map-merge (map1 map2 &optional val-fn)
  (:documentation
    "Returns a map containing all the keys of `map1' and `map2', where the
value for each key contained in only one map is the value from that map, and
the value for each key contained in both maps is the result of calling `val-fn'
on the key, the value from `map1', and the value from `map2'.  `val-fn'
defaults to simply returning its third argument, so the entries in `map2'
simply shadow those in `map1'.  Also, `val-fn' must have the property that if
its second and third arguments are equal, its result is equal to them.  The
default for the new map is computed by calling `val-fn' on the symbol
`fset:map-default' and the defaults for the two maps."))

(defgeneric restrict (map set)
  (:documentation
    "Returns a map containing only those pairs of `map' whose keys are
also in `set'."))

;;; &&& Better name?
(defgeneric restrict-not (map set)
  (:documentation
    "Returns a map containing only those pairs of `map' whose keys are
not in `set'."))

(defgeneric compose (map1 map2-or-fn)
  (:documentation
    "Returns a new map with the same domain as `map1', which maps each member
of that domain to the result of applying first `map1' to it, then applying
`map2-or-fn' to the result."))

(defgeneric first (seq)
  (:documentation
    "Returns the first element of `seq', i.e., element 0.  This has a back-
compatibility method for lists, and adds one for CL sequences generally."))

(defmethod first ((s list))
  (cl:first s))

(defmethod first ((s sequence))
  (elt s 0))

(defgeneric last (seq)
  (:documentation
    "Returns the last element of `seq', i.e., element `(1- (size seq))'.  This
has methods for CL lists and sequences that are NOT COMPATIBLE with `cl:last'.
FSet exports `lastcons' as an arguably better name for the functionality of
`cl:last'."))

(defmethod last ((s list))
  (car (lastcons s)))

(defmethod last ((s sequence))
  (elt s (1- (length s))))

(defgeneric with-first (seq val)
  (:documentation
    "Returns `seq' with `val' prepended, that is, `val' is element 0 of the
result, and the elements of `seq' appear starting at index 1."))

(defgeneric with-last (seq val)
  (:documentation
    "Returns `seq' with `val' appended, that is, `val' is element `(size seq)'
of the result."))

(defgeneric less-first (seq)
  (:documentation
    "Returns the subsequence of `seq' from element 1 through the end."))

(defgeneric less-last (seq)
  (:documentation
    "Returns the subsequence of `seq' from element 0 through the next-to-last
element."))

(defgeneric insert (seq idx val)
  (:documentation
    "Returns a new sequence like `seq' but with `val' inserted at `idx', which
must be in [0, n] where `n' is `(size seq)'."))

;;; &&& Maybe we should shadow `concatenate' instead, so you can specify a
;;; result type.
(defgeneric concat (seq1 seq2)
  (:documentation
    "Returns the concatenation of `seq1' and `seq2'."))

(defgeneric subseq (seq start &optional end)
  (:documentation
    "Returns the subsequence of `seq' from `start' (inclusive) to `end' (exclusive),
where `end' defaults to `(size seq)'."))

(defmethod subseq ((s sequence) start &optional end)
  (cl:subseq s start end))

(defgeneric reverse (seq)
  (:documentation
    "Returns `seq' in reverse order."))

(defmethod reverse ((s sequence))
  (cl:reverse s))

(defgeneric sort (seq pred &key key)
  (:documentation
    "Returns `seq' sorted by `pred', a function of two arguments; if `key' is
supplied, it is a function of one argument that is applied to the elements of
`seq' before they are passed to `pred'.  The sort is not guaranteed to be
stable."))

(defmethod sort ((s sequence) pred &key key)
  (cl:sort s pred :key key))

(defgeneric stable-sort (seq pred &key key)
  (:documentation
    "Returns `seq' sorted by `pred', a function of two arguments; if `key' is
supplied, it is a function of one argument that is applied to the elements of
`seq' before they are passed to `pred'.  The sort is guaranteed to be stable."))

(defmethod stable-sort ((s sequence) pred &key key)
  (cl:stable-sort s pred :key key))


;;; This is the opposite order from `cl:coerce', but I like it better, because I
;;; think the calls are easier to read with the type first.  It's also consistent
;;; with `cl:concatenate' -- the inconsistency between `coerce' and `concatenate'
;;; has long bugged me.
(defgeneric convert (to-type collection &key)
  (:documentation "Converts the collection to the specified type.  Some methods may
take additional keyword arguments to further specify the kind of conversion."))

;;; Generic versions of `find' etc.

(defgeneric find (item collection &key key test)
  (:documentation
    "If `collection' is a Lisp sequence, this simply calls `cl:find'.  On an FSet
collection, the default for `test' is `equal?'; the `:test-not' keyword is not
accepted; and the `:start', `:end', and `:from-end' keywords are accepted only
if `collection' is a seq.  Also, on a map, this scans the domain; on success,
it returns the corresponding range element as the second value."))

(defmethod find (item (s sequence) &rest keyword-args)
  (declare (dynamic-extent keyword-args))
  (apply #'cl:find item s keyword-args))

(defgeneric find-if (pred collection &key key)
  (:documentation
    "If `collection' is a Lisp sequence, this simply calls `cl:find-if'.  On an
FSet collection, the `:start', `:end', and `:from-end' keywords are accepted
only if `collection' is a seq.  Also, on a map, this scans the domain; on
success, it returns the corresponding range element as the second value."))

(defmethod find-if (pred (s sequence) &rest keyword-args)
  (declare (dynamic-extent keyword-args))
  (apply #'cl:find-if pred s keyword-args))

(defgeneric find-if-not (pred collection &key key)
  (:documentation
    "If `collection' is a Lisp sequence, this simply calls `cl:find-if-not'.
On an FSet collection, the `:start', `:end', and `:from-end' keywords are
accepted only if `collection' is a seq.  Also, on a map, this scans the domain;
on success, it returns the corresponding range element as the second value."))

(defmethod find-if-not (pred (s sequence) &rest keyword-args)
  (declare (dynamic-extent keyword-args))
  (apply #'cl:find-if-not pred s keyword-args))

(defgeneric count (item collection &key key test)
  (:documentation
    "If `collection' is a Lisp sequence, this simply calls `cl:count'.  On an FSet
collection, the default for `test' is `equal?'; the `:test-not' keyword is not
accepted; and the `:start', `:end', and `:from-end' keywords are accepted only
if `collection' is a seq.  Also, on a map, this scans the domain."))

(defmethod count (item (s sequence) &rest keyword-args)
  (declare (dynamic-extent keyword-args))
  (apply #'cl:count item s keyword-args))

(defgeneric count-if (pred collection &key key)
  (:documentation
    "If `collection' is a Lisp sequence, this simply calls `cl:count-if'.  On an
FSet collection, the `:start', `:end', and `:from-end' keywords are accepted
only if `collection' is a seq.  Also, on a map, this scans the domain."))

(defmethod count-if (pred (s sequence) &rest keyword-args)
  (declare (dynamic-extent keyword-args))
  (apply #'cl:count-if pred s keyword-args))

(defgeneric count-if-not (pred collection &key key)
  (:documentation
    "If `collection' is a Lisp sequence, this simply calls `cl:count-if-not'.
On an FSet collection, the `:start', `:end', and `:from-end' keywords are
accepted only if `collection' is a seq.  Also, on a map, this scans the domain."))

(defmethod count-if-not (pred (s sequence) &rest keyword-args)
  (declare (dynamic-extent keyword-args))
  (apply #'cl:count-if-not pred s keyword-args))

(defgeneric position (item collection &key key test start end from-end)
  (:documentation
    "If `collection' is a Lisp sequence, this simply calls `cl:position'.  On an
FSet seq, the default for `test' is `equal?', and the `:test-not' keyword is
not accepted."))

(defmethod position (item (s sequence) &rest keyword-args)
  (declare (dynamic-extent keyword-args))
  (apply #'cl:position item s keyword-args))

(defgeneric position-if (pred collection &key key start end from-end)
  (:documentation
    "If `collection' is a Lisp sequence, this simply calls `cl:position-if'.
Also works on an FSet seq."))

(defmethod position-if (pred (s sequence) &rest keyword-args)
  (declare (dynamic-extent keyword-args))
  (apply #'cl:position-if pred s keyword-args))

(defgeneric position-if-not (pred collection &key key)
  (:documentation
    "If `collection' is a Lisp sequence, this simply calls `cl:position-if-not'.
Also works on an FSet seq."))

(defmethod position-if-not (pred (s sequence) &rest keyword-args)
  (declare (dynamic-extent keyword-args))
  (apply #'cl:position-if-not pred s keyword-args))

(defgeneric remove (item collection &key key test start end from-end count)
  (:documentation
    "If `collection' is a Lisp sequence, this simply calls `cl:remove'.  On an
FSet seq, the default for `test' is `equal?', and the `:test-not' keyword is
not accepted."))

(defmethod remove (item (s sequence) &rest keyword-args)
  (declare (dynamic-extent keyword-args))
  (apply #'cl:remove item s keyword-args))

(defgeneric remove-if (pred collection &key key start end from-end count)
  (:documentation
    "If `collection' is a Lisp sequence, this simply calls `cl:remove-if'.
Also works on an FSet seq."))

(defmethod remove-if (pred (s sequence) &rest keyword-args)
  (declare (dynamic-extent keyword-args))
  (apply #'cl:remove-if pred s keyword-args))

(defgeneric remove-if-not (pred collection &key key)
  (:documentation
    "If `collection' is a Lisp sequence, this simply calls `cl:remove-if-not'.
Also works on an FSet seq."))

(defmethod remove-if-not (pred (s sequence) &rest keyword-args)
  (declare (dynamic-extent keyword-args))
  (apply #'cl:remove-if-not pred s keyword-args))

(defgeneric substitute (newitem olditem collection &key key test start end from-end count)
  (:documentation
    "If `collection' is a Lisp sequence, this simply calls `cl:substitute'.  On
an FSet seq, the default for `test' is `equal?', and the `:test-not' keyword
is not accepted."))

(defmethod substitute (newitem olditem (s sequence) &rest keyword-args)
  (declare (dynamic-extent keyword-args))
  (apply #'cl:substitute newitem olditem s keyword-args))

(defgeneric substitute-if (newitem pred collection &key key start end from-end count)
  (:documentation
    "If `collection' is a Lisp sequence, this simply calls `cl:substitute-if'.
Also works on an FSet seq."))

(defmethod substitute-if (newitem pred (s sequence) &rest keyword-args)
  (declare (dynamic-extent keyword-args))
  (apply #'cl:substitute-if newitem pred s keyword-args))

(defgeneric substitute-if-not (newitem pred collection &key key start end from-end count)
  (:documentation
    "If `collection' is a Lisp sequence, this simply calls `cl:substitute-if-not'.
Also works on an FSet seq."))

(defmethod substitute-if-not (newitem pred (s sequence) &rest keyword-args)
  (declare (dynamic-extent keyword-args))
  (apply #'cl:substitute-if-not newitem pred s keyword-args))


;;; ================================================================================
;;; New names for a few existing CL functions

;;; The CL function is poorly (albeit traditionally) named, and we shadow the name.
(defun lastcons (x)
  (cl:last x))

(defun head (x)
  (car x))

(defun tail (x)
  (cdr x))

(declaim (inline lastcons head tail))


;;; ================================================================================
;;; SETF expanders, modify methods, etc.

(define-setf-expander lookup (collection key &environment env)
  "Adds a pair to a map or updates an existing pair, or adds an element to a
sequence or updates an existing element.  This does NOT modify the map or
sequence; it modifies the place (generalized variable) HOLDING the map or
sequence (just like `(setf (ldb ...) ...)').  That is, the `collection' subform
must be `setf'able itself."
  (let ((temps vals stores store-form access-form
	  (get-setf-expansion collection env))
	(key-temp (gensym))
	(val-temp (gensym))
	((coll-temp (car stores))))
    (when (cdr stores)
      (error "Too many values required in `setf' of `lookup'"))
    (values (cons key-temp temps)
	    (cons key vals)
	    (list val-temp)
	    `(let ((,coll-temp (with ,access-form ,key-temp ,val-temp)))
	       ,store-form
	       ,val-temp)
	    `(lookup ,access-form ,key-temp))))

;;; Have to do the same thing for `@' since `setf' would not know what to
;;; do with `@'s normal expansion.
(define-setf-expander @ (collection key &environment env)
  "Adds a pair to a map or updates an existing pair, or adds an element to a
sequence or updates an existing element.  This does NOT modify the map or
sequence; it modifies the place (generalized variable) HOLDING the map or
sequence (just like `(setf (ldb ...) ...)').  That is, the `collection' subform
must be `setf'able itself."
  (let ((temps vals stores store-form access-form
	  (get-setf-expansion collection env))
	(key-temp (gensym))
	(val-temp (gensym))
	((coll-temp (car stores))))
    (when (cdr stores)
      (error "Too many values required in `setf' of `@'"))
    (values (cons key-temp temps)
	    (cons key vals)
	    (list val-temp)
	    `(let ((,coll-temp (with ,access-form ,key-temp ,val-temp)))
	       ,store-form
	       ,val-temp)
	    `(lookup ,access-form ,key-temp))))


(define-modify-macro adjoinf (&rest item-or-tuple)
  with
  "(adjoinf coll . args) --> (setf coll (with coll . args))")

(define-modify-macro removef (&rest item-or-tuple)
  less
  "(removef coll . args) --> (setf coll (less coll . args))")

(define-modify-macro unionf (set)
  union)

(define-modify-macro push-first (val)
  with-first
  "(push-first seq val) --> (setf seq (with-first seq val))")

(define-modify-macro push-last (val)
  with-last
  "(push-last seq val) --> (setf seq (with-last seq val))")

(define-modify-macro pop-first ()
  less-first
  "(pop-first seq) --> (setf seq (less-first seq))")

(define-modify-macro pop-last ()
  less-last
  "(pop-last seq) --> (setf seq (less-last seq))")


;;; ================================================================================
;;; Sets

(defparameter *empty-set* (make-set nil))

(defun empty-set ()
  "Returns an empty set."
  *empty-set*)
(declaim (inline empty-set))

(defmethod empty? ((s set))
  (null (set-contents s)))

(defmethod size ((s set))
  (WB-Set-Tree-Size (set-contents s)))

(defmethod set-size ((s set))
  (WB-Set-Tree-Size (set-contents s)))

(defmethod arb ((s set))
  (let ((tree (set-contents s)))
    (if tree (values (WB-Set-Tree-Arb tree) t)
      (values nil nil))))

(defmethod member? (x (s set))
  (WB-Set-Tree-Member? (set-contents s) x))

(defmethod lookup ((s set) key)
  (WB-Set-Tree-Find-Equal (set-contents s) key))

(defmethod least ((s set))
  (let ((tree (set-contents s)))
    (if tree (values (WB-Set-Tree-Least tree) t)
      (values nil nil))))

(defmethod greatest ((s set))
  (let ((tree (set-contents s)))
    (and tree (values (WB-Set-Tree-Greatest tree) t))))

(defmethod with1 ((s set) value)
  (let ((contents (set-contents s))
	((new-contents (WB-Set-Tree-With contents value))))
    (if (eq new-contents contents)
	s
      (make-set new-contents))))

(defmethod less1 ((s set) value)
  (let ((contents (set-contents s))
	((new-contents (WB-Set-Tree-Less contents value))))
    (if (eq new-contents contents)
	s
      (make-set new-contents))))

(defmethod union ((s1 set) (s2 set))
  (make-set (WB-Set-Tree-Union (set-contents s1) (set-contents s2))))

(defmethod intersection ((s1 set) (s2 set))
  (make-set (WB-Set-Tree-Intersect (set-contents s1) (set-contents s2))))

(defmethod set-difference ((s1 set) (s2 set))
  (make-set (WB-Set-Tree-Diff (set-contents s1) (set-contents s2))))

(defmethod set-difference-2 ((s1 set) (s2 set))
  (let ((newc1 newc2 (WB-Set-Tree-Diff-2 (set-contents s1) (set-contents s2))))
    (values (make-set newc1) (make-set newc2))))

(defmethod subset? ((s1 set) (s2 set))
  (WB-Set-Tree-Subset? (set-contents s1) (set-contents s2)))

(defmethod compare ((s1 set) (s2 set))
  (WB-Set-Tree-Compare (set-contents s1) (set-contents s2)))

(defgeneric internal-do-set (set elt-fn value-fn)
  (:documentation
    "Calls `elt-fn' on successive elements of the set; when done, calls `value-fn'
on no arguments and returns the result(s).  This is called by `do-set' to provide
for the possibility of different set implementations; it is not for public use.
`elt-fn' and `value-fn' must be function objects, not symbols."))

(defmacro do-set ((var set &optional value) &body body)
  "For each member of `set', binds `var' to it and executes `body'.  When done,
returns `value'."
  `(block nil		; in case `body' contains `(return ...)'
     (internal-do-set ,set #'(lambda (,var) . ,body)
			   #'(lambda () ,value))))

(defmethod internal-do-set ((s set) elt-fn value-fn)
  (declare (optimize (speed 3) (safety 0))
	   (type function elt-fn value-fn))
  ;; Expect Python note about "can't use known return convention"
  (Do-WB-Set-Tree-Members (x (set-contents s) (funcall value-fn))
    (funcall elt-fn x)))

(defmethod filter ((pred function) (s set))
  (set-filter pred s))

(defmethod filter ((pred symbol) (s set))
  (set-filter (coerce pred 'function) s))

(defmethod filter ((pred map) (s set))
  (set-filter pred s))

(defmethod filter ((pred set) (s set))
  (intersection pred s))

(defmethod filter ((pred bag) (s set))
  (intersection pred s))

(defun set-filter (pred s)
  (let ((result nil))
    (do-set (x s)
      (when (@ pred x)
	(setq result (WB-Set-Tree-With result x))))
    (make-set result)))

(defmethod image ((fn function) (s set))
  (set-image fn s))

(defmethod image ((fn symbol) (s set))
  (set-image (coerce fn 'function) s))

(defmethod image ((fn map) (s set))
  (set-image fn s))

(defmethod image ((fn set) (s set))
  (set-image fn s))

(defmethod image ((fn bag) (s set))
  (set-image fn s))

(defun set-image (fn s)
  (let ((result nil))
    (do-set (x s)
      (setq result (WB-Set-Tree-With result (@ fn x))))
    (make-set result)))

(defmethod fold ((fn function) (s set) &optional (initial-value nil init?))
  (set-fold fn s initial-value init?))

(defmethod fold ((fn symbol) (s set) &optional (initial-value nil init?))
  (set-fold (coerce fn 'function) s initial-value init?))

(defun set-fold (fn s initial-value init?)
  (declare (optimize (speed 3) (safety 0))
	   (type function fn))
  (if init?
      (let ((result initial-value))
	(do-set (x s)
	  (setq result (funcall fn result x)))
	result)
    (if (empty? s)
	(error "Attempt to fold an empty set with no initial value")
      (let ((result nil)
	    (first? t))
	(do-set (x s)
	  (if first? (setq result x
			   first? nil)
	    (setq result (funcall fn result x))))
	result))))

(defmethod convert ((to-type (eql 'set)) (s set) &key)
  s)

(defmethod convert ((to-type (eql 'bag)) (s set) &key)
  (make-bag (WB-Set-Tree-To-Bag-Tree (set-contents s))))

(defmethod convert ((to-type (eql 'list)) (s set) &key)
  (declare (optimize (speed 3)))
  (let ((result nil))
    (do-set (x s)
      (push x result))
    (nreverse result)))

(defmethod convert ((to-type (eql 'seq)) (s set) &key)
  ;; Not sure we can improve on this much.
  (convert 'seq (convert 'list s)))

(defmethod convert ((to-type (eql 'set)) (l list) &key)
  (make-set (WB-Set-Tree-From-List l)))

(defmethod convert ((to-type (eql 'set)) (s seq) &key)
  (make-set (WB-Seq-Tree-To-Set-Tree (seq-contents s))))

(defmethod find (item (s set) &key key test)
  (declare (optimize (speed 3) (safety 0)))
  (if key
      (let ((key (coerce key 'function)))
	(if test
	    (let ((test (coerce test 'function)))
	      (do-set (x s)
		(when (funcall test item (funcall key x))
		  (return x))))
	  (do-set (x s)
	    (when (equal? item (funcall key x))
	      (return x)))))
    (if (and test (not (or (eq test 'equal?) (eq test #'equal?))))
	(let ((test (coerce test 'function)))
	  (do-set (x s)
	    (when (funcall test item x)
	      (return x))))
      (let ((val? val (lookup item s)))
	(declare (ignore val?))
	val))))

(defmethod find-if (pred (s set) &key key)
  (declare (optimize (speed 3) (safety 0)))
  (let ((pred (coerce pred 'function)))
    (if key
	(let ((key (coerce key 'function)))
	  (do-set (x s)
	    (when (funcall pred (funcall key x))
	      (return x))))
      (do-set (x s)
	(when (funcall pred x)
	  (return x))))))

(defmethod find-if-not (pred (s set) &key key)
  (declare (optimize (speed 3) (safety 0)))
  (let ((pred (coerce pred 'function)))
    (find-if #'(lambda (x) (not (funcall pred x))) s :key key)))

(defmethod count (item (s set) &key key test)
  (declare (optimize (speed 3) (safety 0)))
  (let ((total 0))
    (declare (fixnum total))
    (if key
	(let ((key (coerce key 'function)))
	  (if test
	      (let ((test (coerce test 'function)))
		(do-set (x s total)
		  (when (funcall test item (funcall key x))
		    (incf total))))
	    (do-set (x s total)
	      (when (equal? item (funcall key x))
		(incf total)))))
      (if (and test (not (or (eq test 'equal?) (eq test #'equal?))))
	  (let ((test (coerce test 'function)))
	    (do-set (x s total)
	      (when (funcall test item x)
		(incf total))))
	(let ((val? val (lookup s item)))
	  (declare (ignore val))
	  (if val? 1 0))))))

(defmethod count-if (pred (s set) &key key)
  (declare (optimize (speed 3) (safety 0)))
  (let ((pred (coerce pred 'function))
	(n 0))
    (declare (fixnum n))
    (if key
	(let ((key (coerce key 'function)))
	  (do-set (x s)
	    (when (funcall pred (funcall key x))
	      (incf n))
	    n))
      (do-set (x s)
	(when (funcall pred x)
	  (incf n))
	n))))

(defmethod count-if-not (pred (s set) &key key)
  (declare (optimize (speed 3) (safety 0)))
  (let ((pred (coerce pred 'function)))
    (count-if #'(lambda (x) (not (funcall pred x))) s :key key)))

(defun print-set (set stream level)
  (format stream "#{ ")
  (let ((i 0))
    (do-set (x set)
      (when (> i 0)
	(format stream " "))
      (when (and *print-length* (>= i *print-length*))
	(format stream "...")
	(return))
      (incf i)
      (write x :stream stream :level (and *print-level* (- *print-level* level))))
    (when (> i 0)
      (format stream " ")))
  (format stream "}"))

(gmap::def-gmap-arg-type :set (set)
  `((convert 'list ,set)
    #'null
    #'car
    #'cdr))

(gmap::def-gmap-res-type :set (&optional filterp)
  `(nil #'WB-Set-Tree-With #'make-set ,filterp))


;;; ================================================================================
;;; Bags

(defparameter *empty-bag* (make-bag nil))

(defun empty-bag () *empty-bag*)

(defmethod empty? ((b bag))
  (null (bag-contents b)))

(defmethod arb ((m bag))
  (let ((tree (bag-contents m)))
    (if tree
	(let ((val mult (WB-Bag-Tree-Arb-Pair tree)))
	  (values val mult t))
      (values nil nil nil))))

(defmethod member? (x (b bag))
  (plusp (WB-Bag-Tree-Multiplicity (bag-contents b) x)))

(defmethod lookup ((b bag) x)
  (let ((mult value-found (WB-Bag-Tree-Multiplicity (bag-contents b) x)))
    (if (plusp mult)
	(values t value-found)
      (values nil nil))))

(defmethod least ((b bag))
  (let ((tree (bag-contents b)))
    (if tree
	(let ((val mult (WB-Bag-Tree-Least-Pair tree)))
	  (values val mult t))
      (values nil nil nil))))

(defmethod greatest ((m bag))
  (let ((tree (bag-contents m)))
    (if tree
	(let ((val mult (WB-Bag-Tree-Greatest-Pair tree)))
	  (values val mult t))
      (values nil nil nil))))

(defmethod size ((b bag))
  (WB-Bag-Tree-Total-Count (bag-contents b)))

(defmethod set-size ((b bag))
  (WB-Bag-Tree-Size (bag-contents b)))

(defmethod multiplicity (x (b bag))
  (WB-Bag-Tree-Multiplicity (bag-contents b) x))

(defmethod multiplicity (x (s set))
  (if (member? x s) 1 0))

(defmethod with1 ((b bag) value)
  (make-bag (WB-Bag-Tree-With (bag-contents b) value)))

(defmethod with2 ((b bag) value multiplicity)
  (assert (and (integerp multiplicity) (not (minusp multiplicity))))
  (if (zerop multiplicity) b
    (make-bag (WB-Bag-Tree-With (bag-contents b) value multiplicity))))

(defmethod less1 ((b bag) value)
  (make-bag (WB-Bag-Tree-Less (bag-contents b) value)))

(defmethod less2 ((b bag) value multiplicity)
  (assert (and (integerp multiplicity) (not (minusp multiplicity))))
  (if (zerop multiplicity) b
    (make-bag (WB-Bag-Tree-Less (bag-contents b) value multiplicity))))

(defmethod union ((b1 bag) (b2 bag))
  (make-bag (WB-Bag-Tree-Union (bag-contents b1) (bag-contents b2))))

(defmethod union ((s set) (b bag))
  (make-bag (WB-Bag-Tree-Union (WB-Set-Tree-To-Bag-Tree (set-contents s))
			       (bag-contents b))))

(defmethod union ((b bag) (s set))
  (make-bag (WB-Bag-Tree-Union (bag-contents b)
			       (WB-Set-Tree-To-Bag-Tree (set-contents s)))))

(defmethod bag-sum ((b1 bag) (b2 bag))
  (make-bag (WB-Bag-Tree-Sum (bag-contents b1) (bag-contents b2))))

(defmethod bag-sum ((s set) (b bag))
  (make-bag (WB-Bag-Tree-Sum (WB-Set-Tree-To-Bag-Tree (set-contents s))
			     (bag-contents b))))

(defmethod bag-sum ((b bag) (s set))
  (make-bag (WB-Bag-Tree-Sum (bag-contents b)
			     (WB-Set-Tree-To-Bag-Tree (set-contents s)))))

(defmethod intersection ((s1 bag) (s2 bag))
  (make-bag (WB-Bag-Tree-Intersect (bag-contents s1) (bag-contents s2))))

(defmethod intersection ((s set) (b bag))
  (make-bag (WB-Set-Tree-Intersect (set-contents s)
				   (WB-Bag-Tree-To-Set-Tree (bag-contents b)))))

(defmethod intersection ((b bag) (s set))
  (make-bag (WB-Set-Tree-Intersect (WB-Bag-Tree-To-Set-Tree (bag-contents b))
				   (set-contents s))))

(defmethod bag-product ((b1 bag) (b2 bag))
  (make-bag (WB-Bag-Tree-Product (bag-contents b1) (bag-contents b2))))

(defmethod bag-product ((s set) (b bag))
  (make-bag (WB-Bag-Tree-Product (WB-Set-Tree-To-Bag-Tree (set-contents s))
				 (bag-contents b))))

(defmethod bag-product ((b bag) (s set))
  (make-bag (WB-Bag-Tree-Product (bag-contents b)
				 (WB-Set-Tree-To-Bag-Tree (set-contents s)))))

(defmethod bag-difference ((b1 bag) (b2 bag))
  (make-bag (WB-Bag-Tree-Diff (bag-contents b1) (bag-contents b2))))

(defmethod bag-difference ((s set) (b bag))
  (make-bag (WB-Bag-Tree-Diff (WB-Set-Tree-To-Bag-Tree (set-contents s))
			      (bag-contents b))))

(defmethod bag-difference ((b bag) (s set))
  (make-bag (WB-Bag-Tree-Diff (bag-contents b)
			      (WB-Set-Tree-To-Bag-Tree (set-contents s)))))

(defmethod subbag? ((b1 bag) (b2 bag))
  (WB-Bag-Tree-Subbag? (bag-contents b1) (bag-contents b2)))

(defmethod subbag? ((s set) (b bag))
  (WB-Bag-Tree-Subbag? (WB-Set-Tree-To-Bag-Tree (set-contents s)) (bag-contents b)))

(defmethod subbag? ((b bag) (s set))
  (WB-Bag-Tree-Subbag? (bag-contents b) (WB-Set-Tree-To-Bag-Tree (set-contents s))))

(defmethod compare ((b1 bag) (b2 bag))
  (WB-Bag-Tree-Compare (bag-contents b1) (bag-contents b2)))

(defgeneric internal-do-bag (bag elt-fn value-fn)
  (:documentation
    "Calls `elt-fn' on successive pairs of the bag (the second argument is
the multiplicity); when done, calls `value-fn' on no arguments and returns the
result(s).  This is called by `do-bag' to provide for the possibility of
different bag implementations; it is not for public use.  `elt-fn' and
`value-fn' must be function objects, not symbols."))

(defmacro do-bag ((value-var mult-var bag &optional value)
		  &body body)
  "For each member of `bag', binds `value-var' and `mult-var' to the member and
its multiplicity respectively, and executes `body'.  When done, returns `value'."
  `(block nil
     (internal-do-bag ,bag #'(lambda (,value-var ,mult-var) . ,body)
		      #'(lambda () ,value))))

(defmethod internal-do-bag ((b bag) elt-fn value-fn)
  (declare (optimize (speed 3) (safety 0))
	   (type function elt-fn value-fn))
  ;; Expect Python note about "can't use known return convention"
  (Do-WB-Bag-Tree-Pairs (x n (bag-contents b) (funcall value-fn))
    (funcall elt-fn x n)))

(defmethod filter ((pred function) (b bag))
  (bag-filter pred b))

(defmethod filter ((pred symbol) (b bag))
  (bag-filter (coerce pred 'function) b))

(defmethod filter ((pred map) (b bag))
  (bag-filter pred b))

(defmethod filter ((pred set) (b bag))
  (bag-product (convert pred 'bag) b))

(defmethod filter ((pred bag) (b bag))
  (bag-filter pred b))

(defun bag-filter (pred b)
  (let ((result nil))
    (do-bag (x n b)
      (when (@ pred x)
	(setq result (WB-Bag-Tree-With result x n))))
    (make-bag result)))

(defmethod image ((fn function) (b bag))
  (bag-image fn b))

(defmethod image ((fn symbol) (b bag))
  (bag-image (coerce fn 'function) b))

(defmethod image ((fn map) (b bag))
  (bag-image fn b))

(defmethod image ((fn set) (b bag))
  (bag-image fn b))

(defmethod image ((fn bag) (b bag))
  (bag-image fn b))

(defun bag-image (fn b)
  (let ((result nil))
    (do-bag (x n b)
      (setq result (WB-Bag-Tree-With result (@ fn x) n)))
    (make-bag result)))

(defmethod fold ((fn function) (s bag) &optional (initial-value nil init?))
  (bag-fold fn s initial-value init?))

(defmethod fold ((fn symbol) (s bag) &optional (initial-value nil init?))
  (bag-fold (coerce fn 'function) s initial-value init?))

(defun bag-fold (fn s initial-value init?)
  ;; Expect 5 Python notes about generic arithmetic.
  (declare (optimize (speed 3) (safety 0))
	   (type function fn))
  (if init?
      (let ((result initial-value))
	(do-bag (x n s)
	  (dotimes (i n)
	    (setq result (funcall fn result x))))
	result)
    (if (empty? s)
	(error "Attempt to fold an empty bag with no initial value")
      (let ((result nil)
	    (first? t))
	(do-bag (x n s)
	  (if first? (setq result x
			   first? nil)
	    (setq result (funcall fn result x)))
	  (dotimes (i (1- n))
	    (setq result (funcall fn result x))))
	result))))

(defmethod convert ((to-type (eql 'bag)) (b bag) &key)
  b)

(defmethod convert ((to-type (eql 'list)) (b bag) &key)
  (declare (optimize (speed 3) (safety 0)))
  (let ((result nil))
    (do-bag (value count b)
      ;; Expect 2 Python notes about generic arithmetic.
      (dotimes (i count)
	(push value result)))
    (nreverse result)))

(defmethod convert ((to-type (eql 'alist)) (b bag) &key)
  (declare (optimize (speed 3) (safety 0)))
  (let ((result nil))
    (do-bag (value count b)
      (push (cons value count) result))
    (nreverse result)))

(defmethod convert ((to-type (eql 'bag)) (l list) &key from-type)
  "If `from-type' is the symbol `alist', treats the operand as an alist where the
cdr of each pair (which must be a positive integer) is the member count for
the car.  Otherwise the operand is treated as a simple list of members, some
of which may be repeated."
  (if (eq from-type 'alist)
      (let ((contents nil))
	(dolist (pr l)
	  (unless (and (integerp (cdr pr)) (< 0 (cdr pr)))
	    (error "Cdr of pair is not a positive integer: ~S"
		   pr))
	  (setq contents (WB-Bag-Tree-With contents (car pr) (cdr pr))))
	(make-bag contents))
    ;; &&& Improve me someday
    (let ((contents nil))
      (dolist (x l)
	(setq contents (WB-Bag-Tree-With contents x)))
      (make-bag contents))))

(defmethod find (item (b bag) &key key test)
  (declare (optimize (speed 3) (safety 0)))
  (if key
      (let ((key (coerce key 'function)))
	(if test
	    (let ((test (coerce test 'function)))
	      (do-bag (x n b nil)
		(declare (ignore n))
		(when (funcall test item (funcall key x))
		  (return x))))
	  (do-bag (x n b nil)
		(declare (ignore n))
	    (when (equal? item (funcall key x))
	      (return x)))))
    (if (and test (not (or (eq test 'equal?) (eq test #'equal?))))
	(let ((test (coerce test 'function)))
	  (do-bag (x n b nil)
	    (declare (ignore n))
	    (when (funcall test item x)
	      (return x))))
      (let ((val? val (lookup b item)))
	(declare (ignore val?))
	val))))

(defmethod find-if (pred (b bag) &key key)
  (declare (optimize (speed 3) (safety 0)))
  (let ((pred (coerce pred 'function)))
    (if key
	(let ((key (coerce key 'function)))
	  (do-bag (x n b nil)
	    (declare (ignore n))
	    (when (funcall pred (funcall key x))
	      (return x))))
      (do-bag (x n b nil)
	(declare (ignore n))
	(when (funcall pred x)
	  (return x))))))

(defmethod find-if-not (pred (b bag) &key key)
  (declare (optimize (speed 3) (safety 0)))
  (let ((pred (coerce pred 'function)))
    (find-if #'(lambda (x) (not (funcall pred x))) b :key key)))

(defmethod count (item (b bag) &key key test)
  (declare (optimize (speed 3) (safety 0)))
  (let ((total 0))
    (if key
	(let ((key (coerce key 'function)))
	  (if test
	      (let ((test (coerce test 'function)))
		(do-bag (x n b total)
		  (when (funcall test item (funcall key x))
		    (incf total n))))
	    (do-bag (x n b total)
	      (when (equal? item (funcall key x))
		(incf total n)))))
      (if (and test (not (or (eq test 'equal?) (eq test #'equal?))))
	  (let ((test (coerce test 'function)))
	    (do-bag (x n b total)
	      (when (funcall test item x)
		(incf total n))))
	(multiplicity item b)))))

(defmethod count-if (pred (b bag) &key key)
  (declare (optimize (speed 3) (safety 0)))
  (let ((pred (coerce pred 'function))
	(total 0))
    (if key
	(let ((key (coerce key 'function)))
	  (do-bag (x n b nil)
	    (when (funcall pred (funcall key x))
	      (incf total n))
	    total))
      (do-bag (x n b nil)
	(when (funcall pred x)
	  (incf total n))
	total))))

(defmethod count-if-not (pred (s bag) &key key)
  (declare (optimize (speed 3) (safety 0)))
  (let ((pred (coerce pred 'function)))
    (count-if #'(lambda (x) (not (funcall pred x))) s :key key)))

(defun print-bag (bag stream level)
  (format stream "#{% ")
  (let ((i 0))
    (do-bag (x n bag)
      (when (> i 0)
	(format stream " "))
      (when (and *print-length* (>= i *print-length*))
	(format stream "...")
	(return))
      (incf i)
      (if (> n 1)
	  (progn
	    (format stream "#%")
	    (write `(,x ,n) :stream stream
		   :level (and *print-level* (- *print-level* level))))
	(write x :stream stream :level (and *print-level* (- *print-level* level)))))
    (when (> i 0)
      (format stream " ")))
  (format stream "%}"))


;;; Note that this yields each element potentially multiple times.  (GMap needs
;;; a way for an arg type to return pairs, other than as conses or lists.)
(gmap::def-gmap-arg-type :bag (bag)
  `((convert 'list ,bag)
    #'null
    #'car
    #'cdr))

(gmap::def-gmap-res-type :bag (&optional filterp)
  `(nil #'WB-Bag-Tree-With #'make-bag ,filterp))


;;; ================================================================================
;;; Maps

(defparameter *empty-map* (make-map nil))

(defun empty-map (&optional default)
  (if default (make-map nil default)
    *empty-map*))
(declaim (inline empty-map))

(defmethod empty? ((m map))
  (null (map-contents m)))

(defmethod arb ((m map))
  (let ((tree (map-contents m)))
    (if tree
	(let ((key val (WB-Map-Tree-Arb-Pair tree)))
	  (values key val t))
      (values nil nil nil))))

(defmethod least ((m map))
  (let ((tree (map-contents m)))
    (if tree
	(let ((key val (WB-Map-Tree-Least-Pair tree)))
	  (values key val t))
      (values nil nil nil))))

(defmethod greatest ((m map))
  (let ((tree (map-contents m)))
    (if tree
	(let ((key val (WB-Map-Tree-Greatest-Pair tree)))
	  (values key val t))
      (values nil nil nil))))

(defmethod size ((m map))
  (WB-Map-Tree-Size (map-contents m)))

(defmethod lookup ((m map) key)
  (let ((val? val (WB-Map-Tree-Lookup (map-contents m) key)))
    ;; Our internal convention is the reverse of the external one.
    (values (if val? val (map-default m)) val?)))

(defmethod with2 ((m map) key value)
  (make-map (WB-Map-Tree-With (map-contents m) key value)
	    (map-default m)))

(defmethod less1 ((m map) key)
  (make-map (WB-Map-Tree-Less (map-contents m) key)
	    (map-default m)))

(defmethod domain ((m map))
  ;; &&& Cache this?  It's pretty fast anyway.
  (make-set (WB-Map-Tree-Domain (map-contents m))))

(defmethod compare ((map1 map) (map2 map))
  (WB-Map-Tree-Compare (map-contents map1) (map-contents map2)))

(defgeneric internal-do-map (map elt-fn value-fn)
  (:documentation
    "Calls `elt-fn' on successive pairs of the map (as two arguments); when done,
calls `value-fn' on no arguments and returns the result(s).  This is called by
`do-map' to provide for the possibility of different map implementations; it
is not for public use.  `elt-fn' and `value-fn' must be function objects, not
symbols."))

(defmacro do-map ((key-var value-var map &optional value) &body body)
  "For each pair of `map', binds `key-var' and `value-var' and executes `body'.
When done, returns `value'."
  `(block nil
     (internal-do-map ,map
		      #'(lambda (,key-var ,value-var) . ,body)
		      #'(lambda () ,value))))

(defmethod internal-do-map ((m map) elt-fn value-fn)
  (declare (optimize (speed 3) (safety 0))
	   (type function elt-fn value-fn))
  ;; Expect Python note about "can't use known return convention"
  (Do-WB-Map-Tree-Pairs (x y (map-contents m) (funcall value-fn))
    (funcall elt-fn x y)))

(defmethod filter ((pred function) (m map))
  (map-filter pred m))

(defmethod filter ((pred symbol) (m map))
  (map-filter (coerce pred 'function) m))

(defmethod filter ((pred map) (m map))
  (map-filter pred m))

(defmethod filter ((pred set) (m map))
  (map-filter pred m))

(defmethod filter ((pred bag) (m map))
  (map-filter pred m))

(defun map-filter (pred m)
  (let ((result nil))
    (do-map (x y m)
      (when (@ pred x y)
	(setq result (WB-Map-Tree-With result x y))))
    (make-map result (map-default m))))

(defmethod image ((fn function) (m map))
  (map-image fn m))

(defmethod image ((fn symbol) (m map))
  (map-image (coerce fn 'function) m))

(defun map-image (fn m)
  (let ((result nil))
    (do-map (x y m)
      (let ((new-x new-y (funcall fn x y)))
	(setq result (WB-Map-Tree-With result new-x new-y))))
    (make-map result (map-default m))))

(defmethod range ((m map))
  ;;; &&& Also a candidate for caching -- but the operation isn't terribly common.
  (let ((s nil))
    (do-map (key val m)
      (declare (ignore key))
      (setq s (WB-Set-Tree-With s val)))
    (make-set s)))

(defmethod map-merge ((map1 map) (map2 map)
		      &optional (val-fn #'(lambda (k v1 v2)
					    (declare (ignore k v1))
					    v2)))
  (make-map (WB-Map-Tree-Merge (map-contents map1) (map-contents map2)
			       (coerce val-fn 'function))
	    (funcall val-fn nil (map-default map1) (map-default map2))))

(defmethod restrict ((m map) (s set))
  (make-map (WB-Map-Tree-Restrict (map-contents m) (set-contents s))
	    (map-default m)))

(defmethod restrict-not ((m map) (s set))
  (make-map (WB-Map-Tree-Restrict-Not (map-contents m) (set-contents s))
	    (map-default m)))

(defmethod compose ((map1 map) (map2 map))
  (let ((tree2 (map-contents map2))
	(result nil))
    (do-map (key val1 map1)
      (let ((val2? val2 (WB-Map-Tree-Lookup tree2 val1)))
	(setq result (WB-Map-Tree-With result key (if val2? val2
						    (map-default map2))))))
    (let ((new-default new-default? (WB-Map-Tree-Lookup tree2 (map-default map1))))
      (make-map result (if new-default? new-default (map-default map2))))))

(defmethod compose ((m map) (fn function))
  (map-fn-compose m fn))

(defmethod compose ((m map) (fn symbol))
  (map-fn-compose m (coerce fn 'function)))

(defun map-fn-compose (m fn)
  (declare (optimize (speed 3) (safety 0))
	   (type function fn))
  (let ((result nil))
    (do-map (key val m)
      (setq result (WB-Map-Tree-With result key (funcall fn val))))
    (make-map result (funcall fn (map-default m)))))

(defmethod convert ((to-type (eql 'map)) (m map) &key)
  m)

(defmethod convert ((to-type (eql 'list)) (m map) &key)
  (let ((result nil))
    (do-map (key val m)
      (push (cons key val) result))
    (nreverse result)))

(defmethod convert ((to-type (eql 'seq)) (m map) &key)
  (convert 'seq (convert 'list m)))

(defmethod convert ((to-type (eql 'set)) (m map) &key pair-fn)
  (let ((result nil))
    (do-map (key val m)
      (setq result (WB-Set-Tree-With result (if pair-fn
						(funcall pair-fn key val)
					      (list key val)))))
    (make-set result)))

(defmethod convert ((to-type (eql 'map)) (alist list) &key)
  (let ((m nil))
    (dolist (pr alist)
      (setq m (WB-Map-Tree-With m (car pr) (cdr pr))))
    (make-map m)))

(defmethod find (item (m map) &key key test)
  (declare (optimize (speed 3) (safety 0)))
  (if key
      (let ((key (coerce key 'function)))
	(if test
	    (let ((test (coerce test 'function)))
	      (do-map (x y m nil)
		(when (funcall test item (funcall key x))
		  (return (values x y)))))
	  (do-map (x y m nil)
	    (when (equal? item (funcall key x))
	      (return (values x y))))))
    (if test
	(let ((test (coerce test 'function)))
	  (do-map (x y m nil)
	    (when (funcall test item x)
	      (return (values x y)))))
      (let ((val? val (WB-Map-Tree-Lookup (map-contents m) item)))
	(if val? (values item val)
	  (values nil nil))))))

(defmethod find-if (pred (m map) &key key)
  (declare (optimize (speed 3) (safety 0)))
  (let ((pred (coerce pred 'function)))
    (if key
	(let ((key (coerce key 'function)))
	  (do-map (x y m nil)
	    (when (funcall pred (funcall key x))
	      (return (values x y)))))
      (do-map (x y m nil)
	(when (funcall pred x)
	  (return (values x y)))))))

(defmethod find-if-not (pred (m map) &key key)
  (declare (optimize (speed 3) (safety 0)))
  (let ((pred (coerce pred 'function)))
    (find-if #'(lambda (x) (not (funcall pred x))) m :key key)))

(defmethod count (item (m map) &key key test)
  (declare (optimize (speed 3) (safety 0)))
  (let ((total 0))
    (declare (fixnum total))
    (if key
	(let ((key (coerce key 'function)))
	  (if test
	      (let ((test (coerce test 'function)))
		(do-map (x y m total)
		  (declare (ignore y))
		  (when (funcall test item (funcall key x))
		    (incf total))))
	    (progn
	      (do-map (x y m total)
		(declare (ignore y))
		(when (equal? item (funcall key x))
		  (incf total))))))
      (if (and test (not (or (eq test 'equal?) (eq test #'equal?))))
	  (let ((test (coerce test 'function)))
	    (do-map (x y m total)
	      (declare (ignore y))
	      (when (funcall test item x)
		(incf total))))
	(let ((val? val (WB-Map-Tree-Lookup (map-contents m) item)))
	  (declare (ignore val))
	  (if val? 1 0))))))

(defmethod count-if (pred (m map) &key key)
  (declare (optimize (speed 3) (safety 0)))
  (let ((pred (coerce pred 'function))
	(n 0))
    (declare (fixnum n))
    (if key
	(let ((key (coerce key 'function)))
	  (do-map (x y m)
	    (declare (ignore y))
	    (when (funcall pred (funcall key x))
	      (incf n))
	    n))
      (do-map (x y m)
	(declare (ignore y))
	(when (funcall pred x)
	  (incf n))
	n))))

(defmethod count-if-not (pred (m map) &key key)
  (declare (optimize (speed 3) (safety 0)))
  (let ((pred (coerce pred 'function)))
    (count-if #'(lambda (x) (not (funcall pred x))) m :key key)))

(defun print-map (map stream level)
  (format stream "#{| ")
  (let ((i 0))
    (do-map (x y map)
      (when (> i 0)
	(format stream " "))
      (when (and *print-length* (>= i *print-length*))
	(format stream "...")
	(return))
      (incf i)
      (format stream "(")
      (write x :stream stream :level (and *print-level* (- *print-level* level)))
      (format stream " ")
      (write y :stream stream :level (and *print-level* (- *print-level* level)))
      (format stream ")"))
    (when (> i 0)
      (format stream " ")))
  (format stream "|}")
  (let ((default (map-default map)))
    (when default
      (format stream "/~A" default))))


;;; ================================================================================
;;; Seqs

(defparameter *empty-seq* (make-seq nil))

(defun empty-seq () *empty-seq*)
(declaim (inline empty-seq))

(defmethod empty? ((s seq))
  (null (seq-contents s)))

(defmethod size ((s seq))
  (WB-Seq-Tree-Size (seq-contents s)))

(defmethod lookup ((s seq) key)
  (let ((val? val (WB-Seq-Tree-Subscript (seq-contents s) key)))
    (values val val?)))

(defmethod first ((s seq))
  (let ((val? val (WB-Seq-Tree-Subscript (seq-contents s) 0)))
    (values val val?)))

(defmethod last ((s seq))
  (let ((tree (seq-contents s))
	((val? val (WB-Seq-Tree-Subscript tree (1- (WB-Seq-Tree-Size tree))))))
    (values val val?)))

(defmethod with-first ((s seq) val)
  (make-seq (WB-Seq-Tree-Insert (seq-contents s) 0 val)))

(defmethod with-last ((s seq) val)
  (let ((tree (seq-contents s)))
    (make-seq (WB-Seq-Tree-Insert tree (WB-Seq-Tree-Size tree) val))))

(defmethod less-first ((s seq))
  (let ((tree (seq-contents s)))
    (make-seq (WB-Seq-Tree-Subseq tree 1 (WB-Seq-Tree-Size tree)))))

(defmethod less-last ((s seq))
  (let ((tree (seq-contents s)))
    (make-seq (WB-Seq-Tree-Subseq tree 0 (1- (WB-Seq-Tree-Size tree))))))

(defmethod with2 ((s seq) index val)
  (let ((tree (seq-contents s))
	((size (WB-Seq-Tree-Size tree))))
    (unless (and (>= index 0) (<= index size))
      ;;; &&& Signal a condition?
      (error "Index ~D out of bounds on ~A" index s))
    (make-seq (if (= index size)
		  (WB-Seq-Tree-Insert tree index val)
		(WB-Seq-Tree-With tree index val)))))

(defmethod insert ((s seq) idx val)
  (let ((tree (seq-contents s)))
    (unless (and (>= idx 0) (<= idx (WB-Seq-Tree-Size tree)))
      ;;; &&& Signal a condition?
      (error "Index ~D out of bounds on ~A" idx s))
    (make-seq (WB-Seq-Tree-Insert tree idx val))))

(defmethod less1 ((s seq) idx)
  (let ((tree (seq-contents s)))
    (unless (and (>= idx 0) (< idx (WB-Seq-Tree-Size tree)))
      ;;; &&& Signal a condition?
      (error "Index ~D out of bounds on ~A" idx s))
    (make-seq (WB-Seq-Tree-Remove tree idx))))

(defmethod concat ((s1 seq) (s2 seq))
  (make-seq (WB-Seq-Tree-Concat (seq-contents s1) (seq-contents s2))))

(defmethod subseq ((s seq) start &optional end)
  (let ((tree (seq-contents s)))
    (make-seq (WB-Seq-Tree-Subseq tree start (or end (WB-Seq-Tree-Size tree))))))

(defmethod reverse ((s seq))
  (make-seq (WB-Seq-Tree-Reverse (seq-contents s))))

(defmethod sort ((s seq) pred &key key)
  (convert 'seq (cl:sort (convert 'vector s) pred :key key)))

(defmethod stable-sort ((s seq) pred &key key)
  (convert 'seq (cl:stable-sort (convert 'vector s) pred :key key)))

(defmethod convert ((to-type (eql 'seq)) (s seq) &key)
  s)

(defmethod convert ((to-type (eql 'seq)) (vec vector) &key)
  (make-seq (WB-Seq-Tree-From-Vector vec)))

(defmethod convert ((to-type (eql 'vector)) (s seq) &key)
  (WB-Seq-Tree-To-Vector (seq-contents s)))

(defmethod convert ((to-type (eql 'seq)) (l list) &key)
  (make-seq (WB-Seq-Tree-From-List l)))

(defmethod convert ((to-type (eql 'list)) (s seq) &key)
  (WB-Seq-Tree-To-List (seq-contents s)))

(defmethod compare ((s1 seq) (s2 seq))
  (WB-Seq-Tree-Compare (seq-contents s1) (seq-contents s2)))

(defgeneric internal-do-seq (seq elt-fn value-fn
				 &key start end from-end?)
  (:documentation
    "Calls `elt-fn' on successive elements of `seq', possibly restricted by
`start' and `end', and in reverse order if `from-end?' is true.  When done,
calls `value-fn' on no arguments and returns the result(s).  This is called
by `do-seq' to provide for the possibility of different seq implementations;
it is not for public use.  `elt-fn' and `value-fn' must be function objects,
not symbols."))

(defmacro do-seq ((var seq
		   &key (start nil start?) (end nil end?) (from-end? nil from-end??)
		   (value nil))
		  &body body)
  "For each element of `seq', possibly restricted by `start' and `end', and in
reverse order if `from-end?' is true, binds `var' to it and executes `body'.
When done, returns `value'."
  `(block nil
     (internal-do-seq ,seq
		      #'(lambda (,var) . ,body)
		      #'(lambda () ,value)
		      ,@(and start? `(:start ,start))
		      ,@(and end? `(:end ,end))
		      ,@(and from-end?? `(:from-end? ,from-end?)))))

(defmethod internal-do-seq ((s seq) elt-fn value-fn
			    &key (start 0)
			         (end (WB-Seq-Tree-Size (seq-contents s)))
			         from-end?)
  (declare (optimize (speed 3) (safety 0))
	   (type function elt-fn value-fn))
  ;; Expect Python note about "can't use known return convention"
  (Do-WB-Seq-Tree-Members-Gen (x (seq-contents s) start end from-end?
				 (funcall value-fn))
    (funcall elt-fn x)))

(defmethod member? (x (s seq))
  (declare (optimize (speed 3) (safety 0)))
  (do-seq (y s)
    (when (equal? y x)
      (return t))))

(defmethod filter ((fn function) (s seq))
  (seq-filter fn s))

(defmethod filter ((fn symbol) (s seq))
  (seq-filter (coerce fn 'function) s))

(defmethod filter ((fn map) (s seq))
  (seq-filter #'(lambda (x) (lookup fn x)) s))

(defmethod filter ((fn set) (s seq))
  (seq-filter #'(lambda (x) (lookup fn x)) s))

(defmethod filter ((fn bag) (s seq))
  (seq-filter #'(lambda (x) (lookup fn x)) s))

(defun seq-filter (fn s)
  (declare (optimize (speed 3) (safety 0))
	   (type function fn))
  (let ((result nil))
    (do-seq (x s)
      ;; Since constructing seqs is much faster than for the other types, we
      ;; insist `fn' be a function instead of using `@'.
      (when (funcall fn x)
	(push x result)))
    (make-seq (WB-Seq-Tree-From-List (nreverse result)))))

(defmethod image ((fn function) (s seq))
  (seq-image fn s))

(defmethod image ((fn symbol) (s seq))
  (seq-image (coerce fn 'function) s))

(defmethod image ((fn map) (s seq))
  (seq-image #'(lambda (x) (lookup fn x)) s))

(defmethod image ((fn set) (s seq))
  (seq-image #'(lambda (x) (lookup fn x)) s))

(defmethod image ((fn bag) (s seq))
  (seq-image #'(lambda (x) (lookup fn x)) s))

(defun seq-image (fn s)
  (declare (optimize (speed 3) (safety 0))
	   (type function fn))
  ;; This is not bad, but we could do better by walking the tree of `s' and building
  ;; the result in the same shape.
  (let ((result nil))
    (do-seq (x s)
      ;; Since constructing seqs is much faster than for the other types, we
      ;; insist `fn' be a function instead of using `@'.
      (push (funcall fn x) result))
    (make-seq (WB-Seq-Tree-From-List (nreverse result)))))

(defmethod fold ((fn function) (s seq) &optional (initial-value nil init?))
  (seq-fold fn s initial-value init?))

(defmethod fold ((fn symbol) (s seq) &optional (initial-value nil init?))
  (seq-fold (coerce fn 'function) s initial-value init?))

(defun seq-fold (fn s initial-value init?)
  (declare (optimize (speed 3) (safety 0))
	   (type function fn))
  (if init?
      (let ((result initial-value))
	(do-seq (x s)
	  (setq result (funcall fn result x)))
	result)
    (if (empty? s)
	(error "Attempt to fold an empty sequence with no initial value")
      (let ((result nil)
	    (first? t))
	(do-seq (x s)
	  (if first? (setq result x
			   first? nil)
	    (setq result (funcall fn result x))))
	result))))

(defmethod find (item (s seq) &key key test start end from-end)
  (declare (optimize (speed 3) (safety 0)))
  (if key
      (let ((key (coerce key 'function)))
	(if test
	    (let ((test (coerce test 'function)))
	      (do-seq (x s :start start :end end :from-end? from-end :value nil)
		(when (funcall test item (funcall key x))
		  (return x))))
	  (do-seq (x s :start start :end end :from-end? from-end :value nil)
	    (when (equal? item (funcall key x))
	      (return x)))))
    (if test
	(let ((test (coerce test 'function)))
	  (do-seq (x s :start start :end end :from-end? from-end :value nil)
	    (when (funcall test item x)
		  (return x))))
      (do-seq (x s :start start :end end :from-end? from-end :value nil)
	(when (equal? item x)
	  (return x))))))

(defmethod find-if (pred (s seq) &key key start end from-end)
  (declare (optimize (speed 3) (safety 0)))
  (let ((pred (coerce pred 'function)))
    (if key
	(let ((key (coerce key 'function)))
	  (do-seq (x s :start start :end end :from-end? from-end :value nil)
	    (when (funcall pred (funcall key x))
	      (return x))))
      (do-seq (x s :start start :end end :from-end? from-end :value nil)
	(when (funcall pred x)
	  (return x))))))

(defmethod find-if-not (pred (s seq) &key key start end from-end)
  (declare (optimize (speed 3) (safety 0)))
  (let ((pred (coerce pred 'function)))
    (find-if #'(lambda (x) (not (funcall pred x))) s
	     :key key :start start :end end :from-end from-end)))

(defmethod count (item (s seq) &key key test start end from-end)
  (declare (optimize (speed 3) (safety 0)))
  (let ((total 0))
    (declare (fixnum total))
    (if key
	(let ((key (coerce key 'function)))
	  (if test
	      (let ((test (coerce test 'function)))
		(do-seq (x s :start start :end end :from-end? from-end
			   :value total)
		  (when (funcall test item (funcall key x))
		    (incf total))))
	    (do-seq (x s :start start :end end :from-end? from-end
		       :value total)
	      (when (equal? item (funcall key x))
		(incf total)))))
      (if (and test (not (or (eq test 'equal?) (eq test #'equal?))))
	  (let ((test (coerce test 'function)))
	    (do-seq (x s :start start :end end :from-end? from-end
		       :value total)
	      (when (funcall test item x)
		(incf total))))
	(do-seq (x s :start start :end end :from-end? from-end
		   :value total)
	  (when (equal? item x)
	    (incf total)))))))

(defmethod count-if (pred (s seq) &key key start end from-end)
  (declare (optimize (speed 3) (safety 0)))
  (let ((pred (coerce pred 'function))
	(n 0))
    (declare (fixnum n))
    (if key
	(let ((key (coerce key 'function)))
	  (do-seq (x s :start start :end end :from-end? from-end)
	    (when (funcall pred (funcall key x))
	      (incf n))
	    n))
      (do-seq (x s :start start :end end :from-end? from-end)
	(when (funcall pred x)
	  (incf n))
	n))))

(defmethod count-if-not (pred (s seq) &key key start end from-end)
  (declare (optimize (speed 3) (safety 0)))
  (let ((pred (coerce pred 'function)))
    (count-if #'(lambda (x) (not (funcall pred x))) s
	      :key key :start start :end end :from-end from-end)))

(defmethod position (item (s seq) &key key test start end from-end)
  (declare (optimize (speed 3) (safety 0)))
  (let ((pos 0))
    (declare (fixnum pos))
    (if key
	(let ((key (coerce key 'function)))
	  (if test
	      (let ((test (coerce test 'function)))
		(do-seq (x s :start start :end end :from-end? from-end)
		  (when (funcall test item (funcall key x))
		    (return pos))
		  (incf pos)))
	    (do-seq (x s :start start :end end :from-end? from-end)
	      (when (equal? item (funcall key x))
		(return pos))
	      (incf pos))))
      (if (and test (not (or (eq test 'equal?) (eq test #'equal?))))
	  (let ((test (coerce test 'function)))
	    (do-seq (x s :start start :end end :from-end? from-end)
	      (when (funcall test item x)
		(return pos))
	      (incf pos)))
	(do-seq (x s :start start :end end :from-end? from-end)
	  (when (equal? item x)
	    (return pos))
	  (incf pos))))))

(defmethod position-if (pred (s seq) &key key start end from-end)
  (declare (optimize (speed 3) (safety 0)))
  (let ((pred (coerce pred 'function))
	(pos 0))
    (declare (fixnum pos))
    (if key
	(let ((key (coerce key 'function)))
	  (do-seq (x s :start start :end end :from-end? from-end)
	    (when (funcall pred (funcall key x))
	      (return pos))
	    (incf pos)))
      (do-seq (x s :start start :end end :from-end? from-end)
	(when (funcall pred x)
	  (return pos))
	(incf pos)))))

(defmethod position-if-not (pred (s seq) &key key start end from-end)
  (declare (optimize (speed 3) (safety 0)))
  (let ((pred (coerce pred 'function)))
    (position-if #'(lambda (x) (not (funcall pred x))) s
		 :key key :start start :end end :from-end from-end)))

(defmethod remove (item (s seq) &key key test start end from-end count)
  (declare (optimize (speed 3) (safety 0)))
  (let ((start (or start 0))
	(end (or end (size s)))
	(count (or count (size s)))
	((head (subseq s 0 start))
	 (tail (subseq s end)))
	(mid nil)
	(test (if test (coerce test 'function) #'equal?))
	(key (and key (coerce key 'function))))
    (declare (fixnum count))
    (do-seq (x s :start start :end end :from-end? from-end)
      (if (and (> count 0)
	       (funcall test item (if key (funcall key x) x)))
	  (decf count)
	(push x mid)))
    (concat head (concat (convert 'seq (if from-end mid (nreverse mid)))
			 tail))))

(defmethod remove-if (pred (s seq) &key key start end from-end count)
  (declare (optimize (speed 3) (safety 0)))
  (let ((start (or start 0))
	(end (or end (size s)))
	(count (or count (size s)))
	((head (subseq s 0 start))
	 (tail (subseq s end)))
	(mid nil)
	(pred (coerce pred 'function))
	(key (and key (coerce key 'function))))
    (declare (fixnum count))
    (do-seq (x s :start start :end end :from-end? from-end)
      (if (and (> count 0)
	       (funcall pred (if key (funcall key x) x)))
	  (decf count)
	(push x mid)))
    (concat head (concat (convert 'seq (if from-end mid (nreverse mid)))
			 tail))))

(defmethod remove-if-not (pred (s seq) &key key start end from-end count)
  (declare (optimize (speed 3) (safety 0)))
  (let ((pred (coerce pred 'function)))
    (remove-if #'(lambda (x) (not (funcall pred x))) s
	       :key key :start start :end end :from-end from-end :count count)))

(defmethod substitute (newitem olditem (s seq) &key key test start end from-end count)
  (declare (optimize (speed 3) (safety 0)))
  (let ((start (or start 0))
	(end (or end (size s)))
	(count (or count (size s)))
	((head (subseq s 0 start))
	 (tail (subseq s end)))
	(mid nil)
	(test (if test (coerce test 'function) #'equal?))
	(key (and key (coerce key 'function))))
    (declare (fixnum count))
    (do-seq (x s :start start :end end :from-end? from-end)
      (if (and (> count 0)
	       (funcall test olditem (if key (funcall key x) x)))
	  (progn (push newitem mid) (decf count))
	(push x mid)))
    (concat head (concat (convert 'seq (if from-end mid (nreverse mid)))
			 tail))))

(defmethod substitute-if (newitem pred (s seq) &key key start end from-end count)
  (declare (optimize (speed 3) (safety 0)))
  (let ((start (or start 0))
	(end (or end (size s)))
	(count (or count (size s)))
	((head (subseq s 0 start))
	 (tail (subseq s end)))
	(mid nil)
	(pred (coerce pred 'function))
	(key (and key (coerce key 'function))))
    (declare (fixnum count))
    (do-seq (x s :start start :end end :from-end? from-end)
      (if (and (> count 0)
	       (funcall pred (if key (funcall key x) x)))
	  (progn (push newitem mid) (decf count))
	(push x mid)))
    (concat head (concat (convert 'seq (if from-end mid (nreverse mid)))
			 tail))))

(defmethod substitute-if-not (newitem pred (s seq) &key key start end from-end count)
  (declare (optimize (speed 3) (safety 0)))
  (let ((pred (coerce pred 'function)))
    (substitute-if newitem #'(lambda (x) (not (funcall pred x))) s
		   :key key :start start :end end :from-end from-end :count count)))

(defun print-seq (seq stream level)
  (format stream "#[ ")
  (let ((i 0))
    (do-seq (x seq)
      (when (> i 0)
	(format stream " "))
      (when (and *print-length* (>= i *print-length*))
	(format stream "...")
	(return))
      (incf i)
      (write x :stream stream :level (and *print-level* (- *print-level* level))))
    (when (> i 0)
      (format stream " ")))
  (format stream "]"))

(gmap::def-gmap-arg-type :seq (seq)
  `((convert 'list ,seq)
    #'null
    #'car
    #'cdr))

(gmap::def-gmap-res-type :seq (&optional filterp)
  `(nil #'(lambda (a b) (cons b a)) #'(lambda (s) (convert 'seq (nreverse s)))
    ,filterp))