basis/suffix-arrays/suffix-arrays.factor

   1 ! Copyright (C) 2008 Marc Fauconneau.
   2 ! See http://factorcode.org/license.txt for BSD license.
   3 USING: parser kernel arrays math accessors sequences
   4 math.vectors math.order sorting binary-search sets assocs fry ;
   5 IN: suffix-arrays
   6
   7 <PRIVATE
   8
   9 : suffixes ( string -- suffixes-seq )
  10     dup length [ tail-slice ] with map ;
  11
  12 : prefix<=> ( begin seq -- <=> )
  13     [ <=> ] [ swap head? ] 2bi [ drop +eq+ ] when ;
  14
  15 : find-index ( begin suffix-array -- index/f )
  16     [ prefix<=> ] with search drop ;
  17
  18 : from-to ( index begin suffix-array -- from/f to/f )
  19     swap '[ _ head? not ]
  20     [ find-last-from drop dup [ 1+ ] when ]
  21     [ find-from drop ] 3bi ;
  22
  23 : <funky-slice> ( from/f to/f seq -- slice )
  24     [
  25         tuck
  26         [ drop 0 or ] [ length or ] 2bi*
  27         [ min ] keep
  28     ] keep <slice> ; inline
  29
  30 PRIVATE>
  31
  32 : >suffix-array ( seq -- array )
  33     [ suffixes ] map concat natural-sort ;
  34
  35 : SA{ \ } [ >suffix-array ] parse-literal ; parsing
  36
  37 : query ( begin suffix-array -- matches )
  38     2dup find-index dup
  39     [ -rot [ from-to ] keep <funky-slice> [ seq>> ] map prune ]
  40     [ 3drop { } ] if ;