core/sorting/sorting.factor

   1 ! Copyright (C) 2005, 2008 Slava Pestov.
   2 ! See http://factorcode.org/license.txt for BSD license.
   3 USING: accessors arrays kernel math sequences vectors math.order
   4 sequences sequences.private math.order ;
   5 IN: sorting
   6
   7 ! Optimized merge-sort:
   8 !
   9 ! 1) only allocates 2 temporary arrays
  10
  11 ! 2) first phase (interchanging pairs x[i], x[i+1] where
  12 ! x[i] > x[i+1]) is handled specially
  13
  14 <PRIVATE
  15
  16 TUPLE: merge
  17 { seq    array }
  18 { accum  vector }
  19 { accum1 vector }
  20 { accum2 vector }
  21 { from1  array-capacity }
  22 { to1    array-capacity }
  23 { from2  array-capacity }
  24 { to2    array-capacity } ;
  25
  26 : dump ( from to seq accum -- )
  27     #! Optimize common case where to - from = 1, 2, or 3.
  28     >r >r 2dup swap - r> r> pick 1 =
  29     [ >r >r 2drop r> nth-unsafe r> push ] [
  30         pick 2 = [
  31             >r >r 2drop dup 1+
  32             r> [ nth-unsafe ] curry bi@
  33             r> [ push ] curry bi@
  34         ] [
  35             pick 3 = [
  36                 >r >r 2drop dup 1+ dup 1+
  37                 r> [ nth-unsafe ] curry tri@
  38                 r> [ push ] curry tri@
  39             ] [
  40                 >r nip subseq r> push-all
  41             ] if
  42         ] if
  43     ] if ; inline
  44
  45 : l-elt   [ from1>> ] [ seq>> ] bi nth-unsafe ; inline
  46 : r-elt   [ from2>> ] [ seq>> ] bi nth-unsafe ; inline
  47 : l-done? [ from1>> ] [ to1>> ] bi number= ; inline
  48 : r-done? [ from2>> ] [ to2>> ] bi number= ; inline
  49 : dump-l  [ [ from1>> ] [ to1>> ] [ seq>> ] tri ] [ accum>> ] bi dump ; inline
  50 : dump-r  [ [ from2>> ] [ to2>> ] [ seq>> ] tri ] [ accum>> ] bi dump ; inline
  51 : l-next  [ [ l-elt ] [ [ 1+ ] change-from1 drop ] bi ] [ accum>> ] bi push ; inline
  52 : r-next  [ [ r-elt ] [ [ 1+ ] change-from2 drop ] bi ] [ accum>> ] bi push ; inline
  53 : decide  [ [ l-elt ] [ r-elt ] bi ] dip call +gt+ eq? ; inline
  54
  55 : (merge) ( merge quot: ( elt1 elt2 -- <=> ) -- )
  56     over r-done? [ drop dump-l ] [
  57         over l-done? [ drop dump-r ] [
  58             2dup decide
  59             [ over r-next ] [ over l-next ] if
  60             (merge)
  61         ] if
  62     ] if ; inline recursive
  63
  64 : flip-accum ( merge -- )
  65     dup [ accum>> ] [ accum1>> ] bi eq? [
  66         dup accum1>> underlying>> >>seq
  67         dup accum2>> >>accum
  68     ] [
  69         dup accum1>> >>accum
  70         dup accum2>> underlying>> >>seq
  71     ] if
  72     dup accum>> 0 >>length 2drop ; inline
  73
  74 : <merge> ( seq -- merge )
  75     \ merge new
  76         over >vector >>accum1
  77         swap length <vector> >>accum2
  78         dup accum1>> underlying>> >>seq
  79         dup accum2>> >>accum
  80         dup accum>> 0 >>length drop ; inline
  81
  82 : compute-midpoint ( merge -- merge )
  83     dup [ from1>> ] [ to2>> ] bi + 2/ >>to1 ; inline
  84
  85 : merging ( from to merge -- )
  86     swap >>to2
  87     swap >>from1
  88     compute-midpoint
  89     dup [ to1>> ] [ seq>> length ] bi min >>to1
  90     dup [ to2>> ] [ seq>> length ] bi min >>to2
  91     dup to1>> >>from2
  92     drop ; inline
  93
  94 : nth-chunk ( n size -- from to ) [ * dup ] keep + ; inline
  95
  96 : chunks ( length size -- n ) [ align ] keep /i ; inline
  97
  98 : each-chunk ( length size quot -- )
  99     [ [ chunks ] keep ] dip
 100     [ nth-chunk ] prepose curry
 101     each-integer ; inline
 102
 103 : merge ( from to merge quot -- )
 104     [ [ merging ] keep ] dip (merge) ; inline
 105
 106 : sort-pass ( merge size quot -- )
 107     [
 108         over flip-accum
 109         over [ seq>> length ] 2dip
 110     ] dip
 111     [ merge ] 2curry each-chunk ; inline
 112
 113 : sort-loop ( merge quot -- )
 114     [ 2 [ over seq>> length over > ] ] dip
 115     [ [ 1 shift 2dup ] dip sort-pass ] curry
 116     [ ] while 2drop ; inline
 117
 118 : each-pair ( seq quot -- )
 119     [ [ length 1+ 2/ ] keep ] dip
 120     [ [ 1 shift dup 1+ ] dip ] prepose curry each-integer ; inline
 121
 122 : (sort-pairs) ( i1 i2 seq quot accum -- )
 123     [ 2dup length = ] 2dip rot [
 124         [ drop nip nth ] dip push
 125     ] [
 126         [
 127             [ tuck [ nth-unsafe ] 2bi@ 2dup ] dip call +gt+ eq?
 128             [ swap ] when
 129         ] dip tuck [ push ] 2bi@
 130     ] if ; inline
 131
 132 : sort-pairs ( merge quot -- )
 133     [ [ seq>> ] [ accum>> ] bi ] dip swap
 134     [ (sort-pairs) ] 2curry each-pair ; inline
 135
 136 PRIVATE>
 137
 138 : sort ( seq quot -- sortedseq )
 139     [ <merge> ] dip
 140     [ sort-pairs ] [ sort-loop ] [ drop accum>> underlying>> ] 2tri ;
 141     inline
 142
 143 : natural-sort ( seq -- sortedseq ) [ <=> ] sort ;
 144
 145 : sort-keys ( seq -- sortedseq ) [ [ first ] compare ] sort ;
 146
 147 : sort-values ( seq -- sortedseq ) [ [ second ] compare ] sort ;
 148
 149 : sort-pair ( a b -- c d ) 2dup after? [ swap ] when ;