tree-map: implement recurseAfterMap option

src/packages/tree-map/index.js

@@ -7,25 +7,32 @@ const map = require("../map");
 
 module.exports = function treeMap(tree, mapper, options = {}) {
 	let key = options.key ?? "children";
+	let recurseAfterMap = options.recurseAfterMap ?? false;
+
+	function next(value) {
+		if (isIterable(value)) {
+			return map(value, (item) => step(item));
+		} else if (value != null && typeof value === "object") {
+			return step(value);
+		} else {
+			return value;
+		}
+	}
 
 	function step(subtree) {
+		if (recurseAfterMap) {
 			let mapped = mapper(subtree);
-		let modifiedProperties = {};
+			mapped[key] = next(mapped[key]);
+			return mapped;
+		} else {
+			let input = {
+				... subtree,
+				[key]: next(subtree)
+			};
 
-		if (isIterable(mapped[key])) {
-			modifiedProperties[key] = mapped[key].map((item) => step(item));
-		} else if (mapped[key] != null && typeof mapped[key] === "object") {
-			modifiedProperties[key] = step(mapped[key]);
+			return mapper(input);
 		}
-
-		// We track modified properties separately and (immutably) merge them at the end, because it's the fastest way to ensure that we don't mutate the input object under any circumstances (even after we add multi-key support)
-		return {
-			... mapped,
-			... modifiedProperties
-		};
 	}
 
-	return isIterable(tree)
-		? map(tree, step)
-		: step(tree);
+	return next(tree);
 };