sequences.extras: add start-all and count-subseq for searching/counting all occurences of a subseq

extra/sequences/extras/extras-docs.factor

@@ -188,3 +188,46 @@ HELP: unsurround
     { "seq1" sequence } }
 { $description "Reverses the result of a " { $link surround } " call, stripping off the prefix " { $snippet "seq2" } " and suffix " { $snippet "seq3" } " to restore the original sequence " { $snippet "seq" } "." }
 { $see-also surround } ;
+
+HELP: start-all
+{ $values
+     { "subseq" sequence } { "seq" sequence } { "indices" sequence } }
+{ $description "Outputs the starting indices of the non-overlapping occurences of " { $snippet "subseq" } " in " { $snippet "seq" } "." }
+{ $examples
+    { $example """USING: prettyprint sequences.extras ; "ABA" "ABABA" start-all ."""
+               "{ 0 }"
+    }
+    { $example """USING: prettyprint sequences.extras ; "ABA" "ABAABA" start-all ."""
+      "{ 0 3 }"
+    }
+ } ;
+
+HELP: start-all*
+{ $values
+    { "subseq" sequence } { "seq" sequence } { "indices" sequence } }
+{ $description "Outputs the starting indices of the possibly overlapping occurences of " { $snippet "subseq" } " in " { $snippet "seq" } "." }
+{ $examples
+    { $example """USING: prettyprint sequences.extras ; "ABA" "ABABA" start-all* ."""
+               "{ 0 2 }"
+    } } ;
+
+HELP: count-subseq
+{ $values
+    { "subseq" sequence } { "seq" sequence } { "n" integer } }
+{ $description "Outputs the number of non-overlapping occurences of " { $snippet "subseq" } " in " { $snippet "seq" } "." }
+{ $examples
+    { $example """USING: prettyprint sequences.extras ; "ABA" "ABABA" count-subseq ."""
+               "1"
+    } } ;
+
+
+HELP: count-subseq*
+{ $values
+    { "subseq" sequence } { "seq" sequence } { "n" integer } }
+{ $description "Outputs the number of possibly overlapping occurences of " { $snippet "subseq" } " in " { $snippet "seq" } "." }
+{ $examples
+    { $example """USING: prettyprint sequences.extras ; "ABA" "ABABA" count-subseq* ."""
+               "2"
+    } } ;
+
+{ start-all start-all* count-subseq count-subseq* } related-words

extra/sequences/extras/extras-tests.factor

@@ -174,3 +174,9 @@ IN: sequences.extras.tests
 { V{ 1 3 5 7 9 } } [ 10 iota >vector [ even? ] reject! ] unit-test
 
 { 3/10 } [ 10 iota [ 3 < ] count* ] unit-test
+
+{ { 0 } } [ "ABA" "ABABA" start-all ] unit-test
+{ { 0 2 } } [ "ABA" "ABABA" start-all* ] unit-test
+{ { 0 3 } } [ "ABA" "ABAABA" start-all ] unit-test
+{ 1 } [ "ABA" "ABABA" count-subseq ] unit-test
+{ 2 } [ "ABA" "ABABA" count-subseq* ] unit-test

extra/sequences/extras/extras.factor

@@ -533,3 +533,20 @@ PRIVATE>
 
 : map-find-last-index ( ... seq quot: ( ... elt index -- ... result/f ) -- ... result i elt )
     [ find-last-index ] (map-find-index) ; inline
+
+:: (start-all) ( subseq seq increment -- indices )
+    0
+    [ [ subseq seq ] dip start* dup ]
+    [ [ increment + ] keep ] produce nip ;
+
+: start-all ( subseq seq -- indices )
+    over length (start-all) ; inline
+
+: start-all* ( subseq seq -- indices )
+    1 (start-all) ; inline
+
+: count-subseq ( subseq seq -- n )
+    start-all length ; inline
+
+: count-subseq* ( subseq seq -- n )
+    start-all* length ; inline