From c2b9327b0763f3457fb09db17e22ee9e1e024792 Mon Sep 17 00:00:00 2001
From: Daniel Shahaf <d.s@daniel.shahaf.name>
Date: Wed, 2 Sep 2015 08:28:03 +0000
Subject: [PATCH] Support literal newlines, part 2

---
 highlighters/main/main-highlighter.zsh        | 36 +++++++++----------
 .../main/test-data/vanilla-newline.zsh        | 12 +++----
 2 files changed, 22 insertions(+), 26 deletions(-)

diff --git a/highlighters/main/main-highlighter.zsh b/highlighters/main/main-highlighter.zsh
index f5711c9..ad31b4c 100755
--- a/highlighters/main/main-highlighter.zsh
+++ b/highlighters/main/main-highlighter.zsh
@@ -102,28 +102,24 @@ _zsh_highlight_main_highlighter()
     local substr_color=0
     local style_override=""
     [[ $start_pos -eq 0 && $arg = 'noglob' ]] && highlight_glob=false
-    # advance $start_pos, skipping over whitespace in $buf.
-    ((start_pos+=${#buf[$start_pos+1,-1]}-${#${buf[$start_pos+1,-1]##([[:space:]]|\\[[:space:]])#}}))
-    ((end_pos=$start_pos+${#arg}))
 
-    # Deal with faked colons.
-    #
-    # We can't use the (Z+n+) flag because that elides the end-of-command
-    # token altogether, so 'echo foo\necho bar' (two commands) becomes
-    # indistinguishable from 'echo foo echo bar' (one command with three
-    # words for arguments).
+    # advance $start_pos, skipping over whitespace in $buf.
     if [[ $arg == ';' ]] ; then
-      if [[ $buf[start_pos+1] == ';' ]] ; then
-        # Literally-input colon. That's the normal case.
-      elif [[ $buf[start_pos] == $'\n' ]] ; then
-        # Newline token rendered as a colon.  Empirically we're off-by-one.
-        (( start_pos -= 1 ))
-        (( end_pos -= 1 ))
-      else
-        # Currently, there aren't any other sources of ";" (SEPER) tokens
-        # besides literally-input semicolons and (Z+n+)-converted newlines,
-        # If that ever changes, handle them here.
-      fi
+      # We're looking for either a semicolon or a newline, whichever comes
+      # first.  Both of these are rendered as a ";" (SEPER) by the ${(z)..}
+      # flag.
+      #
+      # We can't use the (Z+n+) flag because that elides the end-of-command
+      # token altogether, so 'echo foo\necho bar' (two commands) becomes
+      # indistinguishable from 'echo foo echo bar' (one command with three
+      # words for arguments).
+      local needle=$'[;\n]'
+      integer offset=${${buf[start_pos+1,-1]}[(i)$needle]}
+      (( start_pos += offset ))
+      (( end_pos += offset ))
+    else
+      ((start_pos+=${#buf[$start_pos+1,-1]}-${#${buf[$start_pos+1,-1]##([[:space:]]|\\[[:space:]])#}}))
+      ((end_pos=$start_pos+${#arg}))
     fi
 
     # Parse the sudo command line
diff --git a/highlighters/main/test-data/vanilla-newline.zsh b/highlighters/main/test-data/vanilla-newline.zsh
index b876928..f0c8eac 100644
--- a/highlighters/main/test-data/vanilla-newline.zsh
+++ b/highlighters/main/test-data/vanilla-newline.zsh
@@ -29,12 +29,12 @@
 # -------------------------------------------------------------------------------------------------
 
 PREBUFFER=$'echo foo; echo bar\n\n\n'
-BUFFER='echo baz; echo qux'
+BUFFER=' echo baz; echo qux'
 
 expected_region_highlight=(
-  "0 4 $ZSH_HIGHLIGHT_STYLES[builtin]" # echo
-  "5 7 $ZSH_HIGHLIGHT_STYLES[default]" # baz
-  "8 9 $ZSH_HIGHLIGHT_STYLES[default]" # semicolon
-  "11 14 $ZSH_HIGHLIGHT_STYLES[builtin]" # echo
-  "16 18 $ZSH_HIGHLIGHT_STYLES[default]" # qux
+  "1 5 $ZSH_HIGHLIGHT_STYLES[builtin]" # echo
+  "6 8 $ZSH_HIGHLIGHT_STYLES[default]" # baz
+  "9 10 $ZSH_HIGHLIGHT_STYLES[default]" # semicolon
+  "12 15 $ZSH_HIGHLIGHT_STYLES[builtin]" # echo
+  "17 19 $ZSH_HIGHLIGHT_STYLES[default]" # qux
 )