core/sorting/sorting.factor

   1 ! Copyright (C) 2005, 2009 Slava Pestov.
   2 ! See http://factorcode.org/license.txt for BSD license.
   3 USING: accessors arrays assocs kernel math math.order sequences
   4 sequences.private vectors ;
   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 : push-unsafe ( elt seq -- )
  27     [ length ] keep
  28     [ set-nth-unsafe ] [ [ 1 + ] dip length<< ] 2bi ; inline
  29
  30 : push-all-unsafe ( from to src dst -- )
  31     [ over - swap ] 2dip [ pick ] dip [ length ] keep
  32     [ [ + ] dip length<< ] 2keep <copy> (copy) drop ; inline
  33
  34 : l-elt ( merge -- elt ) [ from1>> ] [ seq>> ] bi nth-unsafe ; inline
  35
  36 : r-elt ( merge -- elt ) [ from2>> ] [ seq>> ] bi nth-unsafe ; inline
  37
  38 : l-done? ( merge -- ? ) [ from1>> ] [ to1>> ] bi eq? ; inline
  39
  40 : r-done? ( merge -- ? ) [ from2>> ] [ to2>> ] bi eq? ; inline
  41
  42 : dump-l ( merge -- )
  43     [ [ from1>> ] [ to1>> ] [ seq>> ] tri ] [ accum>> ] bi
  44     push-all-unsafe ; inline
  45
  46 : dump-r ( merge -- )
  47     [ [ from2>> ] [ to2>> ] [ seq>> ] tri ] [ accum>> ] bi
  48     push-all-unsafe ; inline
  49
  50 : l-next ( merge -- )
  51     [ [ l-elt ] [ [ 1 + ] change-from1 drop ] bi ] [ accum>> ] bi
  52     push-unsafe ; inline
  53
  54 : r-next ( merge -- )
  55     [ [ r-elt ] [ [ 1 + ] change-from2 drop ] bi ] [ accum>> ] bi
  56     push-unsafe ; inline
  57
  58 : decide ( merge quot: ( elt1 elt2 -- <=> ) -- ? )
  59     [ [ l-elt ] [ r-elt ] bi ] dip call +gt+ eq? ; inline
  60
  61 : (merge) ( merge quot: ( elt1 elt2 -- <=> ) -- )
  62     over r-done? [ drop dump-l ] [
  63         over l-done? [ drop dump-r ] [
  64             2dup decide
  65             [ over r-next ] [ over l-next ] if
  66             (merge)
  67         ] if
  68     ] if ; inline recursive
  69
  70 : flip-accum ( merge -- )
  71     dup [ accum>> ] [ accum1>> ] bi eq? [
  72         dup accum1>> underlying>> >>seq
  73         dup accum2>> >>accum
  74     ] [
  75         dup accum1>> >>accum
  76         dup accum2>> underlying>> >>seq
  77     ] if
  78     dup accum>> 0 >>length 2drop ; inline
  79
  80 : <merge> ( seq -- merge )
  81     \ merge new
  82         over >vector >>accum1
  83         swap length <vector> >>accum2
  84         dup accum1>> underlying>> >>seq
  85         dup accum2>> >>accum ; inline
  86
  87 : compute-midpoint ( merge -- merge )
  88     dup [ from1>> ] [ to2>> ] bi + 2/ >>to1 ; inline
  89
  90 : merging ( from to merge -- )
  91     swap >>to2
  92     swap >>from1
  93     compute-midpoint
  94     dup [ to1>> ] [ seq>> length ] bi min >>to1
  95     dup [ to2>> ] [ seq>> length ] bi min >>to2
  96     dup to1>> >>from2
  97     drop ; inline
  98
  99 : nth-chunk ( n size -- from to ) [ * dup ] keep + ; inline
 100
 101 : chunks ( length size -- n ) [ align ] keep /i ; inline
 102
 103 : each-chunk ( length size quot -- )
 104     [ [ chunks ] keep ] dip
 105     [ nth-chunk ] prepose curry
 106     each-integer ; inline
 107
 108 : merge ( from to merge quot -- )
 109     [ [ merging ] keep ] dip (merge) ; inline
 110
 111 : sort-pass ( merge size quot -- )
 112     [
 113         over flip-accum
 114         over [ seq>> length ] 2dip
 115     ] dip
 116     [ merge ] 2curry each-chunk ; inline
 117
 118 : sort-loop ( merge quot -- )
 119     [ 2 over seq>> length [ over > ] curry ] dip
 120     [ [ 1 shift 2dup ] dip sort-pass ] curry
 121     while 2drop ; inline
 122
 123 : each-pair ( seq quot -- )
 124     [ [ length 1 + 2/ ] keep ] dip
 125     [ [ 1 shift dup 1 + ] dip ] prepose curry each-integer ; inline
 126
 127 : (sort-pairs) ( i1 i2 seq quot accum -- )
 128     [ 2dup length = ] 2dip rot [
 129         [ drop nip nth-unsafe ] dip push-unsafe
 130     ] [
 131         [
 132             [ [ nth-unsafe ] curry bi@ 2dup ] dip call +gt+ eq?
 133             [ swap ] when
 134         ] dip [ push-unsafe ] curry bi@
 135     ] if ; inline
 136
 137 : sort-pairs ( merge quot -- )
 138     [ [ seq>> ] [ accum>> ] bi ] dip swap
 139     [ (sort-pairs) ] 2curry each-pair ; inline
 140
 141 PRIVATE>
 142
 143 : sort ( seq quot: ( obj1 obj2 -- <=> ) -- sortedseq )
 144     [ <merge> ] dip
 145     [ sort-pairs ] [ sort-loop ] [ drop accum>> underlying>> ] 2tri ;
 146     inline
 147
 148 : natural-sort ( seq -- sortedseq ) [ <=> ] sort ;
 149
 150 : sort-with ( seq quot: ( elt -- key ) -- sortedseq )
 151     [ compare ] curry sort ; inline
 152
 153 : inv-sort-with ( seq quot: ( elt -- key ) -- sortedseq )
 154     [ compare invert-comparison ] curry sort ; inline
 155
 156 GENERIC: sort-keys ( obj -- sortedseq )
 157
 158 M: object sort-keys >alist sort-keys ;
 159
 160 M: sequence sort-keys [ first ] sort-with ;
 161
 162 GENERIC: sort-values ( obj -- sortedseq )
 163
 164 M: object sort-values >alist sort-values ;
 165
 166 M: sequence sort-values [ second ] sort-with ;
 167
 168 : sort-pair ( a b -- c d ) 2dup after? [ swap ] when ;