math.statistics: Make a safe version of kth-object. Bounds check k w/ regards to...

diff --git a/basis/math/statistics/statistics-tests.factor b/basis/math/statistics/statistics-tests.factor
index 029c25fec17b73965ba1acc155deade7c7e5059e..f1901d930407c149985786e32b268789c3682f6c 100644 (file)
--- a/basis/math/statistics/statistics-tests.factor
+++ b/basis/math/statistics/statistics-tests.factor
@@ -1,4 +1,5 @@
-USING: assocs kernel math math.functions math.statistics sequences tools.test ;
+USING: assocs kernel math math.functions math.statistics sequences
+math.order tools.test ;
 IN: math.statistics.tests
 
 [ 1 ] [ { 1 } mean ] unit-test
@@ -19,6 +20,10 @@ IN: math.statistics.tests
 { 4 } [ { 1 2 3 4 } 0 kth-largest ] unit-test
 { 2 } [ { 1 2 3 4 } 2 kth-largest ] unit-test
 
+[ { 1 2 3 4 } 30 kth-largest ] [ bounds-error? ] must-fail-with
+[ { 1 2 3 4 } 2 [ [ ] compare ] kth-object ] [ bounds-error? ] must-fail-with
+{ 3 } [ { 1 2 3 4 } 2 [ before? ] kth-object ] unit-test
+
 [ 1 ] [ { 1 } mode ] unit-test
 [ 3 ] [ { 1 2 3 3 3 4 5 6 76 7 2 21 1 3 3 3 } mode ] unit-test
 

diff --git a/basis/math/statistics/statistics.factor b/basis/math/statistics/statistics.factor
index 4c9204bf80cfc840c4c58a0a2a47f74052321926..c79b087e9083ceee203f1c8aae2131b6e21c4f87 100644 (file)
--- a/basis/math/statistics/statistics.factor
+++ b/basis/math/statistics/statistics.factor
@@ -14,9 +14,12 @@ IN: math.statistics
 : harmonic-mean ( seq -- x )
     [ recip ] map-sum recip ;
 
-:: kth-object ( seq k quot: ( x y -- ? ) -- elt )
+<PRIVATE
+
+:: (kth-object) ( seq k quot: ( x y -- ? ) nth-quot exchange-quot -- elt )
     #! Wirth's method, Algorithm's + Data structues = Programs p. 84
     #! The algorithm modifiers seq, so we clone it
+    k seq bounds-check 2drop
     seq clone :> seq
     0 :> i!
     0 :> j!
@@ -30,10 +33,10 @@ IN: math.statistics
         m j!
         [ i j <= ]
         [
-            [ i seq nth-unsafe x quot call ] [ i 1 + i! ] while
-            [ x j seq nth-unsafe quot call ] [ j 1 - j! ] while
+            [ i seq nth-quot call x quot call ] [ i 1 + i! ] while
+            [ x j seq nth-quot call quot call ] [ j 1 - j! ] while
             i j <= [
-                i j seq exchange-unsafe
+                i j seq exchange-quot call
                 i 1 + i!
                 j 1 - j!
             ] when
@@ -44,9 +47,17 @@ IN: math.statistics
     ] while
     k seq nth ; inline
 
-: kth-smallest ( seq k -- elt ) [ < ] kth-object ;
+: kth-object-unsafe ( seq k quot: ( x y -- ? ) -- elt )
+    [ nth-unsafe ] [ exchange-unsafe ] (kth-object) ; inline
+
+PRIVATE>
+
+: kth-object ( seq k quot: ( x y -- ? ) -- elt )
+    [ nth ] [ exchange ] (kth-object) ; inline
+
+: kth-smallest ( seq k -- elt ) [ < ] kth-object-unsafe ;
     
-: kth-largest ( seq k -- elt ) [ > ] kth-object ;
+: kth-largest ( seq k -- elt ) [ > ] kth-object-unsafe ;
 
 : count-relative ( seq k -- lt eq gt )
     [ 0 0 0 ] 2dip '[