extra/multisets/multisets.factor

   1 ! Copyright (C) 2022 Doug Coleman.
   2 ! See http://factorcode.org/license.txt for BSD license.
   3 USING: accessors assocs kernel math parser prettyprint.custom
   4 sequences trees.avl ;
   5 IN: multisets
   6
   7 TUPLE: multiset size avl hash ;
   8
   9 : <multiset> ( -- multiset )
  10     multiset new
  11         0 >>size
  12         <avl> >>avl
  13         H{ } clone >>hash ; inline
  14
  15 : multiset-emplace ( obj multiset -- )
  16     [ dup 1 + ] change-size
  17     [ avl>> set-at ]
  18     [ hash>> swapd push-at ] 3bi ; inline
  19
  20 : multiset-erase ( obj multiset -- )
  21     [
  22         hash>> delete-at* drop
  23     ] [
  24         nip
  25         [ avl>> '[ _ delete-at ] each ]
  26         [ [ length ] dip [ swap - ] change-size drop ] 2bi
  27     ] 2bi ;
  28
  29 : multiset-clear ( multiset -- )
  30     [ hash>> clear-assoc ]
  31     [ avl>> f >>root 0 >>count drop ]
  32     [ 0 >>size drop ] tri ;
  33
  34 : multiset-empty? ( multiset -- ? ) avl>> assoc-size 0 eq? ; inline
  35
  36 : multiset-in? ( multiset obj -- ? ) swap hash>> key? ; inline
  37
  38 : multiset-count ( multiset obj -- n )
  39     swap hash>> at* [ length ] [ drop 0 ] if ; inline
  40
  41 : multiset-members ( multiset -- seq )
  42     avl>> >alist values ; inline
  43
  44 : multiset-each ( multiset quot -- )
  45     [ multiset-members ] dip each ; inline
  46
  47 : >multiset ( seq -- multiset )
  48     <multiset>
  49     [ '[ _ multiset-emplace ] each ] keep ;
  50
  51 SYNTAX: multiset{
  52     \ } [ >multiset ] parse-literal ;
  53
  54 M: multiset pprint-delims drop \ multiset{ \ } ;
  55
  56 M: multiset >pprint-sequence avl>> >alist values ;
  57
  58 M: multiset pprint-narrow? drop t ;
  59
  60 M: multiset pprint* pprint-object ;