From 1b1e8f1da2d99a6c49c3949857c8bee8ef354d18 Mon Sep 17 00:00:00 2001
From: Kaz Kylheku <kaz@kylheku.com>
Date: Wed, 7 Jun 2023 07:04:17 -0700
Subject: New functions keep-keys-if, separate-keys.

* lib.[ch] (keep_keys_if, separate_keys): New functions.

* eval.c (eval_init): keep-keys-if, separate-keys intrinsics
registered.

* txr.1: Documented.

* stdlib/doc-syms.tl: Updated.
---
 eval.c             |   2 +
 lib.c              | 136 ++++++++++++++++++++++++++++++++++++++++++++++++++++-
 lib.h              |   2 +
 stdlib/doc-syms.tl |   2 +
 tests/012/seq.tl   |  15 ++++++
 txr.1              |  52 ++++++++++++++++++++
 6 files changed, 208 insertions(+), 1 deletion(-)

diff --git a/eval.c b/eval.c
index 52829918..8a8c0aed 100644
--- a/eval.c
+++ b/eval.c
@@ -7319,7 +7319,9 @@ void eval_init(void)
   reg_fun(intern(lit("keepql"), user_package), func_n3o(keepql, 2));
   reg_fun(intern(lit("keepqual"), user_package), func_n3o(keepqual, 2));
   reg_fun(intern(lit("keep-if"), user_package), func_n3o(keep_if, 2));
+  reg_fun(intern(lit("keep-keys-if"), user_package), func_n3o(keep_keys_if, 2));
   reg_fun(intern(lit("separate"), user_package), func_n3o(separate, 2));
+  reg_fun(intern(lit("separate-keys"), user_package), func_n3o(separate_keys, 2));
   reg_fun(intern(lit("remq*"), user_package), func_n2(remq_lazy));
   reg_fun(intern(lit("remql*"), user_package), func_n2(remql_lazy));
   reg_fun(intern(lit("remqual*"), user_package), func_n2(remqual_lazy));
diff --git a/lib.c b/lib.c
index 11581b33..4ad63431 100644
--- a/lib.c
+++ b/lib.c
@@ -3141,7 +3141,6 @@ val remove_if(val pred, val seq_in, val keyfun_in)
   }
 }
 
-
 val remq(val obj, val seq, val keyfun)
 {
   return rem_impl(eq, lit("remq"), obj, seq, keyfun);
@@ -3177,6 +3176,71 @@ val keep_if(val pred, val seq, val keyfun)
   return remove_if(notf(pred), seq, keyfun);
 }
 
+val keep_keys_if(val pred, val seq_in, val keyfun_in)
+{
+  val self = lit("keep-keys-if");
+  val keyfun = default_null_arg(keyfun_in);
+
+  switch (type(seq_in)) {
+  case NIL:
+    return nil;
+  case CONS:
+  case LCONS:
+  case COBJ:
+    {
+      list_collect_decl (out, ptail);
+      val list = seq_in;
+
+      gc_hint(list);
+
+      for (; list; list = cdr(list)) {
+        val elem = car(list);
+        val key = keyfun ? funcall1(keyfun, elem) : elem;
+
+        if (funcall1(pred, key))
+          ptail = list_collect(ptail, key);
+      }
+      return out;
+    }
+  case LIT:
+  case STR:
+  case LSTR:
+    {
+      val out = mkustring(zero);
+      val str = seq_in;
+      cnum len = c_fixnum(length_str(str), self), i;
+
+      for (i = 0; i < len; i++) {
+        val elem = chr_str(str, num_fast(i));
+        val key = keyfun ? funcall1(keyfun, elem) : elem;
+
+        if (funcall1(pred, key))
+          string_extend(out, key, tnil(i == len - 1));
+      }
+
+      return out;
+    }
+  case VEC:
+    {
+      val out = vector(zero, nil);
+      val vec = seq_in;
+      cnum len = c_fixnum(length_vec(vec), self), i;
+
+      for (i = 0; i < len; i++) {
+        val elem = vecref(vec, num_fast(i));
+        val key = keyfun ? funcall1(keyfun, elem) : elem;
+
+        if (funcall1(pred, key))
+          vec_push(out, key);
+      }
+
+      return out;
+    }
+  default:
+    uw_throwf(error_s, lit("~a: ~s isn't a sequence"), self, seq_in, nao);
+  }
+}
+
 val separate(val pred, val seq_in, val keyfun_in)
 {
   val self = lit("separate");
@@ -3260,6 +3324,76 @@ val separate(val pred, val seq_in, val keyfun_in)
   }
 }
 
+val separate_keys(val pred, val seq_in, val keyfun_in)
+{
+  val self = lit("separate-keys");
+  val keyfun = default_null_arg(keyfun_in);
+
+  switch (type(seq_in)) {
+  case NIL:
+    return cons(nil, cons(nil, nil));
+  case CONS:
+  case LCONS:
+  case COBJ:
+    {
+      list_collect_decl (yea, yptail);
+      list_collect_decl (nay, nptail);
+      val list = seq_in;
+
+      gc_hint(list);
+
+      for (; list; list = cdr(list)) {
+        val elem = car(list);
+        val key = keyfun ? funcall1(keyfun, elem) : elem;
+        val is_yea = if3(funcall1(pred, key), t, nil);
+
+        if (is_yea)
+          yptail = list_collect(yptail, key);
+        else
+          nptail = list_collect(nptail, key);
+      }
+
+      return cons(yea, cons(nay, nil));
+    }
+  case LIT:
+  case STR:
+  case LSTR:
+    {
+      val yea = mkustring(zero);
+      val nay = mkustring(zero);
+      val str = seq_in;
+      cnum len = c_fixnum(length_str(str), self), i;
+
+      for (i = 0; i < len; i++) {
+        val elem = chr_str(str, num_fast(i));
+        val key = keyfun ? funcall1(keyfun, elem) : elem;
+
+        string_extend(funcall1(pred, key) ? yea : nay, key, tnil(i == len - 1));
+      }
+
+      return cons(yea, cons(nay, nil));
+    }
+  case VEC:
+    {
+      val yea = vector(zero, nil);
+      val nay = vector(zero, nil);
+      val vec = seq_in;
+      cnum len = c_fixnum(length_vec(vec), self), i;
+
+      for (i = 0; i < len; i++) {
+        val elem = vecref(vec, num_fast(i));
+        val key = keyfun ? funcall1(keyfun, elem) : elem;
+
+        vec_push(funcall1(pred, key) ? yea : nay, key);
+      }
+
+      return cons(yea, cons(nay, nil));
+    }
+  default:
+    uw_throwf(error_s, lit("~a: ~s isn't a sequence"), self, seq_in, nao);
+  }
+}
+
 static val rem_lazy_rec(val obj, val list, val env, val func);
 
 static val rem_lazy_func(val env, val lcons)
diff --git a/lib.h b/lib.h
index 8448db55..1046092f 100644
--- a/lib.h
+++ b/lib.h
@@ -847,7 +847,9 @@ val keepq(val obj, val seq, val keyfun);
 val keepql(val obj, val seq, val keyfun);
 val keepqual(val obj, val seq, val keyfun);
 val keep_if(val pred, val seq, val keyfun);
+val keep_keys_if(val pred, val seq_in, val keyfun_in);
 val separate(val pred, val seq, val keyfun);
+val separate_keys(val pred, val seq_in, val keyfun_in);
 val remq_lazy(val obj, val list);
 val remql_lazy(val obj, val list);
 val remqual_lazy(val obj, val list);
diff --git a/stdlib/doc-syms.tl b/stdlib/doc-syms.tl
index 94896888..89731a67 100644
--- a/stdlib/doc-syms.tl
+++ b/stdlib/doc-syms.tl
@@ -1079,6 +1079,7 @@
       ("juxt" "N-0106CD7F")
       ("keep-if" "N-0159C541")
       ("keep-if*" "N-0159C541")
+      ("keep-keys-if" "N-03D17450")
       ("keep-match-products" "N-01A846D2")
       ("keep-matches" "N-01A846D2")
       ("keepq" "N-00583609")
@@ -1715,6 +1716,7 @@
       ("select" "N-031D7F72")
       ("self-path" "N-03561A65")
       ("separate" "N-0159C541")
+      ("separate-keys" "N-03D17450")
       ("seq-begin" "N-0068A845")
       ("seq-next" "N-02E3D643")
       ("seq-reset" "N-01CA6912")
diff --git a/tests/012/seq.tl b/tests/012/seq.tl
index a65d8ffa..4d2260e9 100644
--- a/tests/012/seq.tl
+++ b/tests/012/seq.tl
@@ -550,3 +550,18 @@
   (search-all '#"x x x x x" '#"y") nil
   (search-all '#"x x x x x" '#"x") (0 1 2 3 4)
   (search-all '#"x x x" "") (0 1 2 3))
+
+(mtest
+  [keep-keys-if evenp (range 1 20) square] (4 16 36 64 100 144 196 256 324 400)
+  [keep-keys-if chr-isupper "foo bar" chr-toupper] "FOOBAR"
+  [keep-keys-if evenp (vec-list (range 1 20)) square] #(4 16 36 64 100 144 196 256 324 400))
+
+
+(mtest
+  [separate-keys evenp (range 1 20) square] ((4 16 36 64 100 144 196 256 324 400)
+                                             (1 9 25 49 81 121 169 225 289 361))
+  [separate-keys chr-isupper "foo bar" chr-toupper] ("FOOBAR" " ")
+  [separate-keys evenp (vec-list (range 1 20)) square] (#(4 16 36 64 100 144 196 256 324 400)
+                                                        #(1 9 25 49 81 121 169 225 289 361)))
+
+
diff --git a/txr.1 b/txr.1
index 434fefea..347be419 100644
--- a/txr.1
+++ b/txr.1
@@ -35285,6 +35285,58 @@ but produce lazy lists.
   -> (("defg" 5))
 .brev
 
+.coNP Functions @ keep-keys-if and @ separate-keys
+.synb
+.mets (keep-keys-if < predicate-fun < sequence <> [ key-fun ])
+.mets (separate-keys < predicate-fun < sequence <> [ key-fun ])
+.syne
+.desc
+The functions
+.code keep-keys-if
+and
+.code separate-keys
+are derived, respectively, from the functions
+.code keep-if
+and
+.codn separate ,
+and have the same syntax and argument semantics. They differ in that
+rather than accumulating the elements of the input
+.codn sequence ,
+they accumulate the transformed values of those elements, as projected
+through the
+.metn key-fun .
+
+Thus when
+.meta key-fun
+is omitted, thus defaulting to
+.codn identity ,
+or else explicitly specified as
+.code identity
+or equivalent function, the behavior of these functions is the
+almost the same as that of
+.code keep-if
+and
+.codn separate .
+However, there may be a difference in whether the output shares structure with
+.metn sequence .
+
+.TP* Example:
+
+.verb
+  ;; square the values 1 to 20, keeping the even squares
+  [keep-keys-if evenp (range 1 20) square]
+  -> (4 16 36 64 100 144 196 256 324 400)
+
+  ;; square the values 1 to 20 separating into even and odd:
+  [separate-keys evenp (range 1 20) square]
+  -> ((4 16 36 64 100 144 196 256 324 400)
+      (1 9 25 49 81 121 169 225 289 361))
+
+  ;; contrast with keep-if: values are of input sequence
+  [keep-if evenp (range 1 20) square]
+  -> (2 4 6 8 10 12 14 16 18 20)
+.brev
+
 .coNP Functions @, countqual @ countql and @ countq
 .synb
 .mets (countq < object << iterable )
-- 
cgit v1.2.3