counted? test

squint-cljs · Dec 9, 2023 · 1bf46f3 · 1bf46f3
1 parent 7d9e2f8
commit 1bf46f3
Show file tree

Hide file tree

Showing 3 changed files with 35 additions and 1 deletion.
diff --git a/resources/squint/core.edn b/resources/squint/core.edn
@@ -69,6 +69,7 @@
   filterv
   find
   first
+  flatten
   fn_QMARK_
   fnil
   frequencies
@@ -156,7 +157,9 @@
   second
   select_keys
   seq
+  seq_QMARK_
   seqable_QMARK_
+  sequential_QMARK_
   set
   shuffle
   some
@@ -178,6 +181,7 @@
   take_while
   to_array
   transduce
+  tree_seq
   true_QMARK_
   truth_
   update

diff --git a/src/squint/core.js b/src/squint/core.js
@@ -457,6 +457,12 @@ export function get(coll, key, otherwise = undefined) {
   return v !== undefined ? v : otherwise;
 }
 
+export function seq_QMARK_(x) {
+  return x != null && !!x[Symbol.iterator];
+}
+
+export const sequential_QMARK_ = seq_QMARK_;
+
 export function seqable_QMARK_(x) {
   return (
     x === null ||
@@ -1915,7 +1921,7 @@ export function boolean_QMARK_(x) {
 export function counted_QMARK_(x) {
   const tc = typeConst(x);
   switch (tc) {
-    case (ARRAY_TYPE, MAP_TYPE, OBJECT_TYPE, LIST_TYPE, SET_TYPE):
+    case ARRAY_TYPE: case MAP_TYPE: case OBJECT_TYPE: case LIST_TYPE: case SET_TYPE:
       return true;
   }
   return false;
@@ -2040,3 +2046,21 @@ export function not_empty(x) {
     return x;
   } else return null;
 }
+
+export function tree_seq(isBranch, children, root) {
+  const walk = function*(node) {
+    yield node;
+    if (truth_(isBranch(node))) {
+      for (const c of children(node)) {
+        yield* walk(c);
+      }
+    }
+  };
+  return lazy(function* () {
+    yield* walk(root);
+  });
+}
+
+export function flatten(x) {
+  return filter(complement(sequential_QMARK_), rest(tree_seq(sequential_QMARK_, seq, x)));
+}
diff --git a/test/squint/compiler_test.cljs b/test/squint/compiler_test.cljs
@@ -2103,5 +2103,11 @@ new Foo();")
                    f))]
     (is (eq [1 1 1 1] (f js-obj)))))
 
+(deftest flatten-test
+  (is (eq [1 2 3 4 3] (jsv! '(vec (flatten '(1 2 (3 (4 (((3))))))))))))
+
+(deftest counted?-test
+  (is (true? (jsv! '(counted? {})))))
+
 (defn init []
   (t/run-tests 'squint.compiler-test 'squint.jsx-test 'squint.string-test))