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