basis/suffix-arrays/suffix-arrays.factor

   1 ! Copyright (C) 2008 Marc Fauconneau.
   2 ! See http://factorcode.org/license.txt for BSD license.
   3 USING: accessors binary-search fry kernel math math.order parser
   4 sequences sets sorting ;
   5 IN: suffix-arrays
   6
   7 <PRIVATE
   8
   9 : suffixes ( string -- suffixes-seq )
  10     dup length <iota> [ 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 : query-from ( index begin suffix-array -- from )
  19     swap '[ _ head? not ] find-last-from drop [ 1 + ] [ 0 ] if* ;
  20
  21 : query-to ( index begin suffix-array -- to )
  22     [ swap '[ _ head? not ] find-from drop ] [ length or ] bi ;
  23
  24 : query-range ( index begin suffix-array -- from to )
  25     [ query-from ] [ query-to ] 3bi [ min ] keep ;
  26
  27 : (query) ( index begin suffix-array -- matches )
  28     [ query-range ] keep <slice> [ seq>> ] map members ;
  29
  30 PRIVATE>
  31
  32 : >suffix-array ( seq -- suffix-array )
  33     members [ suffixes ] map concat natural-sort ;
  34
  35 SYNTAX: SA{ \ } [ >suffix-array ] parse-literal ;
  36
  37 : query ( begin suffix-array -- matches )
  38     [ find-index ] 2keep '[ _ _ (query) ] [ { } ] if* ;