From 20c092b41cb9e9e3ae81b31e7b4b3739fce39af0 Mon Sep 17 00:00:00 2001
From: Vic Nightfall <vic@nightfall.moe>
Date: Sun, 3 Dec 2023 16:52:35 +0100
Subject: [PATCH] Fix underlines a bit

---
 src/parser.pr         | 21 +++++++++++++++++----
 test/test_compiler.pr | 32 ++++++++++++++++----------------
 2 files changed, 33 insertions(+), 20 deletions(-)

diff --git a/src/parser.pr b/src/parser.pr
index dd0516a1..3dbc8f83 100644
--- a/src/parser.pr
+++ b/src/parser.pr
@@ -1428,7 +1428,19 @@ export def identifier_to_str(node: &Node, types: bool = true) -> Str {
 
 def make_node(kind: NodeKind, line: int, column: int, state: &ParseState) -> &Node {
     state.last_string = null
-    let current_token = peek(state)
+
+    var end_line = 0
+    var end_column = 0
+
+    if state.last_token {
+        end_line = state.last_token.value.end_line
+        end_column = state.last_token.value.end_column
+    } else {
+        let current_token = peek(state)
+        end_line = current_token.line
+        end_column = current_token.column
+    }
+
     return {
         kind = kind,
         loc = {
@@ -1437,8 +1449,8 @@ def make_node(kind: NodeKind, line: int, column: int, state: &ParseState) -> &No
             state.module,
             line,
             column,
-            current_token.line,
-            current_token.column
+            end_line,
+            end_column
         } !SourceLoc
     } !Node
 }
@@ -2690,9 +2702,10 @@ def parse_post_expression(parse_state: &ParseState) -> &Node {
             node = expect_func_args(parse_state, node)
         } else if next_token(parse_state, lexer::TokenType::O_SQUARE) {
             skip_newline(parse_state)
-            node = make_bin_op(parse_state, token, NodeKind::ARRAY_SUBSCRIPT, node, parse_expression(parse_state))
+            let expr = parse_expression(parse_state)
             skip_newline(parse_state)
             expect(parse_state, lexer::TokenType::C_SQUARE, "Expected ']'")
+            node = make_bin_op(parse_state, token, NodeKind::ARRAY_SUBSCRIPT, node, expr)
         } else if next_token(parse_state, lexer::TokenType::DOT) {
             token = peek(parse_state)
             var right: &Node
diff --git a/test/test_compiler.pr b/test/test_compiler.pr
index 0adbdb16..f6c4cbc4 100644
--- a/test/test_compiler.pr
+++ b/test/test_compiler.pr
@@ -250,17 +250,17 @@ def #test test_vardecl_fail {
         |
         |main@4:13
         |        var c = 10
-        |            ^~
+        |            ^
         |Redeclaration of `c`
         |
         |main@5:16
         |        var e, (c) = 10
-        |               ^~~~
+        |               ^~~
         |Must assign a value
         |
         |main@6:13
         |        var f: int = "string"
-        |            ^~~~~~~
+        |            ^~~~~~
         |Incompatible types [char] and int\n""")
 }
 
@@ -280,22 +280,22 @@ def #test test_var_let {
     assert env.err() == strip_margin("""
         |main@3:9
         |        a = null
-        |        ^~
+        |        ^
         |Assignment to non var
         |
         |main@4:9
         |        @a = 20
-        |        ^~~
+        |        ^~
         |Assignment to non var
         |
         |main@7:9
         |        b[0] = 20
-        |        ^~~
+        |        ^~~~
         |Assignment to non var
         |
         |main@10:9
         |        c[0] = 20
-        |        ^~~
+        |        ^~~~
         |Assignment to non var\n""")
 }
 
@@ -362,12 +362,12 @@ def #test test_assign_fail {
     assert env.err() == strip_margin("
         |main@4:9
         |        foo = 40
-        |        ^~~~
+        |        ^~~
         |Assignment to non var
         |
         |main@6:9
         |        bar = \"string\"
-        |        ^~~~
+        |        ^~~
         |Incompatible types, can't assign [char] to int
         |
         |main@2:13
@@ -457,7 +457,7 @@ def #test test_def_fail {
     assert env.err() == strip_margin("""
         |main@3:13
         |        def foo {}
-        |            ^~~~
+        |            ^~~
         |Function `foo` was already declared previously (same arguments)
         |
         |main@17:22
@@ -477,7 +477,7 @@ def #test test_def_fail {
         |
         |main@7:24
         |            export def nested_function {}
-        |                       ^~~~~~~~~~~~~~~~
+        |                       ^~~~~~~~~~~~~~~
         |Invalid modifier to closure
         |
         |main@10:22
@@ -558,7 +558,7 @@ def #test test_typedecl_fail {
         |
         |main@7:14
         |        type module::A = int
-        |             ^~~~~~~~~~
+        |             ^~~~~~~~~
         |Can't declare type in sub scope\n""")
 }
 
@@ -714,7 +714,7 @@ def #test test_import_fail {
     assert env.err() == strip_margin("
         |main@2:16
         |        import test::a as A
-        |               ^~~~~~~~
+        |               ^~~~~~~
         |Module `test::a` could not be found
         |
         |main@3:9
@@ -909,7 +909,7 @@ def #test test_array_inference_fail {
     assert env.err() == strip_margin("
         |main@2:13
         |        let a: [?; float] = [1, 2, 3, 4]
-        |            ^~~~~~~~~~~~~~
+        |            ^~~~~~~~~~~~~
         |Incompatible types [4; int] and [?; float]\n")
 }
 
@@ -954,7 +954,7 @@ def #test test_call_fail {
         |
         |main@7:9
         |        a = no_return()
-        |        ^~
+        |        ^
         |Incompatible types, can't assign void to int\n")
 }
 
@@ -1014,7 +1014,7 @@ def #test test_closure_fail {
     assert env.err() == strip_margin("
         |main@3:24
         |            export def inner_function {}
-        |                       ^~~~~~~~~~~~~~~
+        |                       ^~~~~~~~~~~~~~
         |Invalid modifier to closure
         |
         |main@4:32