feat(stdlib): Optimize List.init to support large lists (#2249)

spotandjake · web-flow · commit 3de64ba88cb3 · 2025-02-25T19:18:38.000Z
diff --git a/compiler/test/stdlib/list.test.gr b/compiler/test/stdlib/list.test.gr
@@ -265,7 +265,13 @@ assert sub(0, 2, listSub) == [1, 2]
 assert sub(1, 2, listSub) == [2, 3]
 
 // List.init
+let largeListSize = 500_000
+let mut largeList = []
+for (let mut i = largeListSize - 1; i >= 0; i -= 1) {
+  largeList = [i, ...largeList]
+}
 
+assert init(largeListSize, i => i) == largeList
 assert init(4, idx => idx) == [0, 1, 2, 3]
 assert init(0, idx => fail "Shouldn't be called") == []
 assert init(-10, idx => fail "Shouldn't be called") == []
diff --git a/stdlib/list.gr b/stdlib/list.gr
@@ -13,6 +13,24 @@ from "runtime/unsafe/memory" include Memory
 from "runtime/unsafe/wasmi32" include WasmI32
 from "runtime/dataStructures" include DataStructures
 
+/**
+ * Creates a new list with all elements in reverse order.
+ *
+ * @param list: The list to reverse
+ * @returns The new list
+ *
+ * @since v0.1.0
+ */
+provide let reverse = list => {
+  let rec revHelp = (list, acc) => {
+    match (list) {
+      [] => acc,
+      [first, ...rest] => revHelp(rest, [first, ...acc]),
+    }
+  }
+  revHelp(list, [])
+}
+
 /**
  * Creates a new list of the specified length where each element is
  * initialized with the result of an initializer function. The initializer
@@ -27,17 +45,18 @@ from "runtime/dataStructures" include DataStructures
  * @since v0.3.0
  */
 provide let init = (length, fn) => {
-  // This method can be further improved by checking the length against a specific size
-  // and determining if it can be made tail-recursive, then use List.reverse on it.
-  // Which would be the same as OCaml does it in https://github.com/ocaml/ocaml/blob/03839754f46319aa36d9dad56940a6f3c3bcb48a/stdlib/list.ml#L79
-  let rec iter = (idx, length) => {
-    if (idx >= length) {
-      []
-    } else {
-      [fn(idx), ...iter(idx + 1, length)]
+  if (length < 2048) {
+    // For small lists, use a faster non-tail-recursive version
+    let rec iter = (idx, length) => {
+      if (idx >= length) [] else [fn(idx), ...iter(idx + 1, length)]
+    }
+    iter(0, length)
+  } else {
+    let rec iter = (idx, length, acc) => {
+      if (idx >= length) acc else iter(idx + 1, length, [fn(idx), ...acc])
     }
+    reverse(iter(0, length, []))
   }
-  iter(0, length)
 }
 
 /**
@@ -74,24 +93,6 @@ provide let isEmpty = list => {
   }
 }
 
-/**
- * Creates a new list with all elements in reverse order.
- *
- * @param list: The list to reverse
- * @returns The new list
- *
- * @since v0.1.0
- */
-provide let reverse = list => {
-  let rec iter = (list, acc) => {
-    match (list) {
-      [] => acc,
-      [first, ...rest] => iter(rest, [first, ...acc]),
-    }
-  }
-  iter(list, [])
-}
-
 /**
  * Creates a new list with the elements of the first list followed by
  * the elements of the second list.
diff --git a/stdlib/list.md b/stdlib/list.md
@@ -25,6 +25,31 @@ from "list" include List
 
 Functions and constants included in the List module.
 
+### List.**reverse**
+
+<details disabled>
+<summary tabindex="-1">Added in <code>0.1.0</code></summary>
+No other changes yet.
+</details>
+
+```grain
+reverse : (list: List<a>) => List<a>
+```
+
+Creates a new list with all elements in reverse order.
+
+Parameters:
+
+|param|type|description|
+|-----|----|-----------|
+|`list`|`List<a>`|The list to reverse|
+
+Returns:
+
+|type|description|
+|----|-----------|
+|`List<a>`|The new list|
+
 ### List.**init**
 
 <details disabled>
@@ -116,31 +141,6 @@ Returns:
 |----|-----------|
 |`Bool`|`true` if the list is empty and `false` otherwise|
 
-### List.**reverse**
-
-<details disabled>
-<summary tabindex="-1">Added in <code>0.1.0</code></summary>
-No other changes yet.
-</details>
-
-```grain
-reverse : (list: List<a>) => List<a>
-```
-
-Creates a new list with all elements in reverse order.
-
-Parameters:
-
-|param|type|description|
-|-----|----|-----------|
-|`list`|`List<a>`|The list to reverse|
-
-Returns:
-
-|type|description|
-|----|-----------|
-|`List<a>`|The new list|
-
 ### List.**append**
 
 <details disabled>
