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 ;
   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 ;
  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 M: bit-set delete
  23     ! This isn't allowed to throw an error if the elt wasn't
  24     ! in the set
  25     over integer? [
  26         table>> 2dup bounds-check? [
  27             [ f ] 2dip set-nth
  28         ] [ 2drop ] if
  29     ] [ 2drop ] if ;
  30
  31 ! If you do binary set operations with a bit-set, it's expected
  32 ! that the other thing can also be represented as a bit-set
  33 ! of the same length.
  34 <PRIVATE
  35
  36 ERROR: check-bit-set-failed ;
  37
  38 : check-bit-set ( bit-set -- bit-set )
  39     dup bit-set? [ check-bit-set-failed ] unless ; inline
  40
  41 : bit-set-map ( seq1 seq2 quot -- seq )
  42     [ 2drop length>> ]
  43     [
  44         [
  45             [ [ length ] bi@ assert= ]
  46             [ [ underlying>> ] bi@ ] 2bi
  47         ] dip 2map
  48     ] 3bi bit-array boa ; inline
  49
  50 : (bit-set-op) ( set1 set2 -- table1 table2 )
  51     [ set-like ] keep [ table>> ] bi@ ; inline
  52
  53 : bit-set-op ( set1 set2 quot: ( a b -- c ) -- bit-set )
  54     [ (bit-set-op) ] dip bit-set-map bit-set boa ; inline
  55
  56 PRIVATE>
  57
  58 M: bit-set union
  59     [ bitor ] bit-set-op ;
  60
  61 M: bit-set intersect
  62     [ bitand ] bit-set-op ;
  63
  64 M: bit-set diff
  65     [ bitnot bitand ] bit-set-op ;
  66
  67 M: bit-set subset?
  68     [ intersect ] keep = ;
  69
  70 M: bit-set members
  71     [ table>> length iota ] keep [ in? ] curry filter ;
  72
  73 <PRIVATE
  74
  75 : bit-set-like ( set bit-set -- bit-set' )
  76     ! Throws an error if there are keys that can't be put
  77     ! in the bit set
  78     over bit-set? [ 2dup [ table>> length ] bi@ = ] [ f ] if
  79     [ drop ] [
  80         [ members ] dip table>> length <bit-set>
  81         [ [ adjoin ] curry each ] keep
  82     ] if ;
  83
  84 PRIVATE>
  85
  86 M: bit-set set-like
  87     bit-set-like check-bit-set ; inline
  88
  89 M: bit-set clone
  90     table>> clone bit-set boa ;
  91
  92 M: bit-set cardinality
  93     table>> bit-count ;