diff --git a/basis/generalizations/generalizations-docs.factor b/basis/generalizations/generalizations-docs.factor index e3a7c2d7e4..de74dd1ead 100644 --- a/basis/generalizations/generalizations-docs.factor +++ b/basis/generalizations/generalizations-docs.factor @@ -50,6 +50,11 @@ HELP: firstn } } ; +HELP: set-firstn +{ $values { "n" integer } } +{ $description "A generalization of " { $link set-first } " " +"that sets the first " { $snippet "n" } " elements of a sequence from the top " { $snippet "n" } " elements of the stack." } ; + HELP: npick { $values { "n" integer } } { $description "A generalization of " { $link dup } ", " @@ -257,7 +262,7 @@ HELP: nweave HELP: n*quot { $values { "n" integer } { "quot" quotation } - { "quot'" quotation } + { "quotquot" quotation } } { $examples { $example "USING: generalizations prettyprint math ;" @@ -314,6 +319,7 @@ ARTICLE: "sequence-generalizations" "Generalized sequence operations" narray nsequence firstn + set-firstn nappend nappend-as } ; diff --git a/basis/generalizations/generalizations-tests.factor b/basis/generalizations/generalizations-tests.factor index f95ba63228..d466d56251 100644 --- a/basis/generalizations/generalizations-tests.factor +++ b/basis/generalizations/generalizations-tests.factor @@ -40,6 +40,8 @@ IN: generalizations.tests [ { "xyc" "xyd" } ] [ "x" "y" { "c" "d" } [ 3append ] 2 nwith map ] unit-test [ 1 2 3 4 ] [ { 1 2 3 4 } 4 firstn ] unit-test +[ { 1 2 3 4 } ] [ 1 2 3 4 { f f f f } [ 4 set-firstn ] keep ] unit-test +[ 1 2 3 4 { f f f } [ 4 set-firstn ] keep ] must-fail [ ] [ { } 0 firstn ] unit-test [ "a" ] [ { "a" } 1 firstn ] unit-test diff --git a/basis/generalizations/generalizations.factor b/basis/generalizations/generalizations.factor index 5ca00018a2..2e9d560ae6 100644 --- a/basis/generalizations/generalizations.factor +++ b/basis/generalizations/generalizations.factor @@ -48,6 +48,18 @@ MACRO: nrot ( n -- ) MACRO: -nrot ( n -- ) 1 - [ ] [ '[ swap _ dip ] ] repeat ; +MACRO: set-firstn-unsafe ( n -- ) + [ 1 + ] + [ iota [ '[ _ rot [ set-nth-unsafe ] keep ] ] map ] bi + '[ _ -nrot _ spread drop ] ; + +MACRO: set-firstn ( n -- ) + dup zero? [ drop [ drop ] ] [ + [ 1 - swap bounds-check 2drop ] + [ set-firstn-unsafe ] + bi-curry '[ _ _ bi ] + ] if ; + MACRO: ndrop ( n -- ) [ drop ] n*quot ;