From 8cd34535a829ba4e3207e52ba933147adac47a5c Mon Sep 17 00:00:00 2001
From: Dmitry Vasilev <dmitry.vasil@gmail.com>
Date: Fri, 13 Oct 2023 19:14:43 +0300
Subject: [PATCH] fix coloring

---
 src/color.js | 46 +++++++++++++++++++++++++++++++++++++++++++++-
 test/test.js | 24 ++++++++++++++++++++++++
 2 files changed, 69 insertions(+), 1 deletion(-)

diff --git a/src/color.js b/src/color.js
index 2b4a0ea..e49d75b 100644
--- a/src/color.js
+++ b/src/color.js
@@ -1,3 +1,21 @@
+const collect_function_exprs = new Function('root', `
+  const result = []
+  function do_collect(n) {
+    if(n.type == 'function_expr') {
+      result.push(n)
+      return
+    }
+    if(n.children == null) {
+      return
+    }
+    for(let c of n.children) {
+      do_collect(c)
+    }
+  }
+  do_collect(root)
+  return result
+`)
+
 const is_result_eq = (a,b) => a.result == null
   ? b.result == null
   : b.result != null 
@@ -158,7 +176,33 @@ const do_color = (node, is_root = false) => {
   } 
 
   if(node.result?.error != null) {
-    return [node_to_color(node)]
+    const color = node_to_color(node)
+    const exprs = collect_function_exprs(node)
+    if(exprs.length == 0) {
+      return [color]
+    }
+    // Color node in red, but make holes for function exprs
+    return exprs
+      .map((e, i) => {
+        let prev_index
+        if(i == 0) {
+          prev_index = node.index
+        } else {
+          const prev_expr = exprs[i - 1]
+          prev_index = prev_expr.index + prev_expr.length
+        }
+        return {
+          index: prev_index, 
+          length: e.index - prev_index, 
+          result: color.result
+        }
+      })
+      .concat([{
+        index: exprs.at(-1).index + exprs.at(-1).length, 
+        length: node.index + node.length 
+          - (exprs.at(-1).index + exprs.at(-1).length), 
+        result: color.result,
+      }])
   } 
 
   const result = color_children(node, is_root)
diff --git a/test/test.js b/test/test.js
index f514bdd..ae3c0dc 100644
--- a/test/test.js
+++ b/test/test.js
@@ -1662,6 +1662,30 @@ const y = x()`
     assert_equal(color_body.result.ok, true)
   }),
 
+  test('coloring error with nested fns', () => {
+    const code = `[1].map(_ => {throw new Error()}).map(x => x + 1)`
+    const i = test_initial_state(code)
+    const coloring = color_file(i, '')
+
+    const result = {ok: false, error_origin: true}
+    assert_equal(
+      coloring,
+      [
+        {
+          index: 0, 
+          length: code.indexOf('_ =>'),
+          result
+        },
+        {
+          index: code.indexOf(').map(x =>'), 
+          length: 1,
+          result
+        },
+      ]
+      
+    )
+  }),
+
   test('better parse errors', () => {
     const code = `
       const x = z => {