basis/bit-sets/bit-sets.factor

   1 ! Copyright (C) 2009 Slava Pestov.
   2 ! See http://factorcode.org/license.txt for BSD license.
   3 USING: kernel accessors sequences byte-arrays bit-arrays math
   4 math.bitwise hints sets sequences.private ;
   5 IN: bit-sets
   6
   7 TUPLE: bit-set { table bit-array read-only } ;
   8
   9 : <bit-set> ( capacity -- bit-set )
  10     <bit-array> bit-set boa ; inline
  11
  12 INSTANCE: bit-set set
  13
  14 M: bit-set in?
  15     over integer? [ table>> ?nth ] [ 2drop f ] if ; inline
  16
  17 M: bit-set adjoin
  18     ! This is allowed to throw an error when the elt couldn't
  19     ! go in the set
  20     [ t ] 2dip table>> set-nth ;
  21
  22 <PRIVATE
  23
  24 : ?set-nth ( elt n seq -- )
  25     2dup bounds-check? [ set-nth-unsafe ] [ 3drop ] if ; inline
  26
  27 PRIVATE>
  28
  29 M: bit-set delete
  30     ! This isn't allowed to throw an error if the elt wasn't
  31     ! in the set
  32     over integer? [
  33         [ f ] 2dip table>> ?set-nth
  34     ] [ 2drop ] if ;
  35
  36 ! If you do binary set operations with a bit-set, it's expected
  37 ! that the other thing can also be represented as a bit-set
  38 ! of the same length.
  39 <PRIVATE
  40
  41 ERROR: check-bit-set-failed ;
  42
  43 : check-bit-set ( bit-set -- bit-set )
  44     dup bit-set? [ check-bit-set-failed ] unless ; inline
  45
  46 : bit-set-map ( seq1 seq2 quot -- seq )
  47     [ 2drop length>> ]
  48     [
  49         [
  50             [ [ length ] bi@ assert= ]
  51             [ [ underlying>> ] bi@ ] 2bi
  52         ] dip 2map
  53     ] 3bi bit-array boa ; inline
  54
  55 : (bit-set-op) ( set1 set2 -- table1 table2 )
  56     [ set-like ] keep [ table>> ] bi@ ; inline
  57
  58 : bit-set-op ( set1 set2 quot: ( a b -- c ) -- bit-set )
  59     [ (bit-set-op) ] dip bit-set-map bit-set boa ; inline
  60
  61 PRIVATE>
  62
  63 M: bit-set union
  64     [ bitor ] bit-set-op ;
  65
  66 M: bit-set intersect
  67     [ bitand ] bit-set-op ;
  68
  69 M: bit-set diff
  70     [ bitnot bitand ] bit-set-op ;
  71
  72 M: bit-set subset?
  73     [ intersect ] keep = ;
  74
  75 M: bit-set members
  76     [ table>> length iota ] keep [ in? ] curry filter ;
  77
  78 <PRIVATE
  79
  80 : bit-set-like ( set bit-set -- bit-set' )
  81     ! Throws an error if there are keys that can't be put
  82     ! in the bit set
  83     over bit-set? [ 2dup [ table>> length ] same? ] [ f ] if
  84     [ drop ] [
  85         [ members ] dip table>> length <bit-set>
  86         [ [ adjoin ] curry each ] keep
  87     ] if ;
  88
  89 PRIVATE>
  90
  91 M: bit-set set-like
  92     bit-set-like check-bit-set ; inline
  93
  94 M: bit-set clone
  95     table>> clone bit-set boa ;
  96
  97 M: bit-set cardinality
  98     table>> bit-count ;