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         [ drop 0 or ] [ length or ] bi-curry bi*
  26         [ min ] keep
  27     ] keep <slice> ; inline
  28
  29 PRIVATE>
  30
  31 : >suffix-array ( seq -- array )
  32     [ suffixes ] map concat natural-sort ;
  33
  34 SYNTAX: SA{ \ } [ >suffix-array ] parse-literal ;
  35
  36 : query ( begin suffix-array -- matches )
  37     2dup find-index dup
  38     [ -rot [ from-to ] keep <funky-slice> [ seq>> ] map prune ]
  39     [ 3drop { } ] if ;