From a238647df9a51d33f55585f8ceebe8d839fcdcb3 Mon Sep 17 00:00:00 2001
From: Matthew Martin <phy1729@gmail.com>
Date: Wed, 18 Mar 2020 23:31:11 -0500
Subject: [PATCH] main: Add arithmetic substitution highlighting

Closes #607 #649 #704
---
 changelog.md                                  |   9 +-
 docs/highlighters/main.md                     |   1 +
 highlighters/main/main-highlighter.zsh        | 115 ++++++++++++++++--
 .../main/test-data/arith-cmdsubst-mess.zsh    |  20 +--
 highlighters/main/test-data/arith1.zsh        |   2 +-
 highlighters/main/test-data/arith2.zsh        |   1 +
 6 files changed, 121 insertions(+), 27 deletions(-)

diff --git a/changelog.md b/changelog.md
index d62ad15..1dd1c1d 100644
--- a/changelog.md
+++ b/changelog.md
@@ -50,13 +50,6 @@
 - Fix `echo >&p` highlighting the `p` as a filename if a file by that name happened to exist
   [part of #645]
 
-- Fix `: $((42))` being highlighted as a subshell.
-  [part of #607]
-
-- Regress highlighting of `: $((ls); (ls))`: is a subshell, but will now be
-  incorrectly highlighted as an arithmetic expansion.
-  [#704]
-
 - Fix wrong highlighting of unquoted parameter expansions under zsh 5.2 and older
   [e165f18c758e]
 
@@ -89,6 +82,8 @@
   (such as `;`, `|`, `&&`) before a newline
   [#677; had regressed in 0.7.0]
 
+- Highlight arithmetic expansions (e.g., `$(( 42 ))`)
+  [#607 #649 #704]
 
 # Changes in version 0.7.1
 
diff --git a/docs/highlighters/main.md b/docs/highlighters/main.md
index 5eec335..cc6186b 100644
--- a/docs/highlighters/main.md
+++ b/docs/highlighters/main.md
@@ -42,6 +42,7 @@ This highlighter defines the following styles:
 * `command-substitution-delimiter-quoted` - a quoted command substitution delimiters (`"$(` and `)"`)
 * `process-substitution` - process substitutions (`<(echo foo)`)
 * `process-substitution-delimiter` - process substitution delimiters (`<(` and `)`)
+* `arithmetic-expansion` - arithmetic expansion `$(( 42 ))`)
 * `single-hyphen-option` - single-hyphen options (`-o`)
 * `double-hyphen-option` - double-hyphen options (`--option`)
 * `back-quoted-argument` - backtick command substitution (`` `foo` ``)
diff --git a/highlighters/main/main-highlighter.zsh b/highlighters/main/main-highlighter.zsh
index 9bf0e3a..6b3610c 100644
--- a/highlighters/main/main-highlighter.zsh
+++ b/highlighters/main/main-highlighter.zsh
@@ -1350,8 +1350,13 @@ _zsh_highlight_main_highlighter_highlight_argument()
           (( i = REPLY ))
           highlights+=($reply)
           continue
-        elif [[ $arg[i+1] == $'\x28' && ${arg:$i} != $'\x28\x28'*$'\x29\x29'* ]]; then
-          # command substitution that doesn't look like an arithmetic expansion
+        elif [[ $arg[i+1] == $'\x28' ]]; then
+          if [[ $arg[i+2] == $'\x28' ]] && _zsh_highlight_main_highlighter_highlight_arithmetic $i; then
+            # Arithmetic expansion
+            (( i = REPLY ))
+            highlights+=($reply)
+            continue
+          fi
           start=$i
           (( i += 2 ))
           _zsh_highlight_main_highlighter_highlight_list $(( start_pos + i - 1 )) S $has_end $arg[i,-1]
@@ -1366,10 +1371,6 @@ _zsh_highlight_main_highlighter_highlight_argument()
             highlights+=($(( start_pos + i - 1)) $(( start_pos + i )) command-substitution-delimiter-unquoted)
           fi
           continue
-        else
-          # TODO: if it's an arithmetic expansion, skip past it, to prevent
-          # multiplications from being highlighted as globbing (issue #607,
-          # test-data/arith1.zsh)
         fi
         while [[ $arg[i+1] == [=~#+'^'] ]]; do
           (( i += 1 ))
@@ -1497,11 +1498,17 @@ _zsh_highlight_main_highlighter_highlight_double_quote()
               # $#, $*, $@, $?, $- - like $$ above
               (( k += 1 )) # highlight both dollar signs
               (( i += 1 )) # don't consider the second one as introducing another parameter expansion
-            elif [[ $arg[i+1] == $'\x28' && ${arg:$i} != $'\x28\x28'*$'\x29\x29'* ]]; then
-              # command substitution that doesn't look like an arithmetic expansion
+            elif [[ $arg[i+1] == $'\x28' ]]; then
+              saved_reply=($reply)
+              if [[ $arg[i+2] == $'\x28' ]] && _zsh_highlight_main_highlighter_highlight_arithmetic $i; then
+                # Arithmetic expansion
+                (( i = REPLY ))
+                reply=($saved_reply $reply)
+                continue
+              fi
+
               breaks+=( $last_break $(( start_pos + i - 1 )) )
               (( i += 2 ))
-              saved_reply=($reply)
               _zsh_highlight_main_highlighter_highlight_list $(( start_pos + i - 1 )) S $has_end $arg[i,-1]
               ret=$?
               (( i += REPLY ))
@@ -1682,6 +1689,96 @@ _zsh_highlight_main_highlighter_highlight_backtick()
   REPLY=$i
 }
 
+# Highlight special chars inside arithmetic expansions
+_zsh_highlight_main_highlighter_highlight_arithmetic()
+{
+  local -a saved_reply
+  local style
+  integer i j k paren_depth ret
+  reply=()
+
+  for (( i = $1 + 3 ; i <= end_pos - start_pos ; i += 1 )) ; do
+    (( j = i + start_pos - 1 ))
+    (( k = j + 1 ))
+    case "$arg[$i]" in
+      [\'\"\\@{}])
+        style=unknown-token
+        ;;
+      '(')
+        (( paren_depth++ ))
+        continue
+        ;;
+      ')')
+        if (( paren_depth )); then
+          (( paren_depth-- ))
+          continue
+        fi
+        [[ $arg[i+1] == ')' ]] && { (( i++ )); break; }
+        # Special case ) at the end of the buffer to avoid flashing command substitution for a character
+        (( has_end && (len == k) )) && break
+        # This is a single paren and there are no open parens, so this isn't an arithmetic expansion
+        return 1
+        ;;
+      '`')
+        saved_reply=($reply)
+        _zsh_highlight_main_highlighter_highlight_backtick $i
+        (( i = REPLY ))
+        reply=($saved_reply $reply)
+        continue
+        ;;
+      '$' )
+        if [[ $arg[i+1] == $'\x28' ]]; then
+          saved_reply=($reply)
+          if [[ $arg[i+2] == $'\x28' ]] && _zsh_highlight_main_highlighter_highlight_arithmetic $i; then
+            # Arithmetic expansion
+            (( i = REPLY ))
+            reply=($saved_reply $reply)
+            continue
+          fi
+
+          (( i += 2 ))
+          _zsh_highlight_main_highlighter_highlight_list $(( start_pos + i - 1 )) S $has_end $arg[i,end_pos]
+          ret=$?
+          (( i += REPLY ))
+          reply=(
+            $saved_reply
+            $j $(( start_pos + i )) command-substitution-quoted
+            $j $(( j + 2 )) command-substitution-delimiter-quoted
+            $reply
+          )
+          if (( ret == 0 )); then
+            reply+=($(( start_pos + i - 1 )) $(( start_pos + i )) command-substitution-delimiter)
+          fi
+          continue
+        else
+          continue
+        fi
+        ;;
+      ($histchars[1]) # ! - may be a history expansion
+        if [[ $arg[i+1] != ('='|$'\x28'|$'\x7b'|[[:blank:]]) ]]; then
+          style=history-expansion
+        else
+          continue
+        fi
+        ;;
+      *)
+        continue
+        ;;
+
+    esac
+    reply+=($j $k $style)
+  done
+
+  if [[ $arg[i] != ')' ]]; then
+    # If unclosed, i points past the end
+    (( i-- ))
+  fi
+    style=arithmetic-expansion
+  reply=($(( start_pos + $1 - 1)) $(( start_pos + i )) arithmetic-expansion $reply)
+  REPLY=$i
+}
+
+
 # Called with a single positional argument.
 # Perform filename expansion (tilde expansion) on the argument and set $REPLY to the expanded value.
 #
diff --git a/highlighters/main/test-data/arith-cmdsubst-mess.zsh b/highlighters/main/test-data/arith-cmdsubst-mess.zsh
index 6f60469..82268ac 100644
--- a/highlighters/main/test-data/arith-cmdsubst-mess.zsh
+++ b/highlighters/main/test-data/arith-cmdsubst-mess.zsh
@@ -33,14 +33,14 @@ BUFFER=$': $((ls); (ls))'
 expected_region_highlight=(
   '1 1 builtin' # :
   '3 15 default' # $((ls); (ls))
-  '3 15 command-substitution-unquoted "issue #704"' # $((ls); (ls))
-  '3 4 command-substitution-delimiter-unquoted "issue #704"' # $(
-  '5 5 reserved-word "issue #704"' # (
-  '6 7 command "issue #704"' # ls
-  '8 8 reserved-word "issue #704"' # )
-  '9 9 commandseparator "issue #704"' # ;
-  '11 11 reserved-word "issue #704"' # (
-  '12 13 command "issue #704"' # ls
-  '14 14 reserved-word "issue #704"' # )
-  '15 15 command-substitution-delimiter-unquoted "issue #704"' # )
+  '3 15 command-substitution-unquoted' # $((ls); (ls))
+  '3 4 command-substitution-delimiter-unquoted' # $(
+  '5 5 reserved-word' # (
+  '6 7 command' # ls
+  '8 8 reserved-word' # )
+  '9 9 commandseparator' # ;
+  '11 11 reserved-word' # (
+  '12 13 command' # ls
+  '14 14 reserved-word' # )
+  '15 15 command-substitution-delimiter-unquoted' # )
 )
diff --git a/highlighters/main/test-data/arith1.zsh b/highlighters/main/test-data/arith1.zsh
index 92fa3da..0462f73 100644
--- a/highlighters/main/test-data/arith1.zsh
+++ b/highlighters/main/test-data/arith1.zsh
@@ -33,5 +33,5 @@ BUFFER=$': $(( 6 * 9 ))'
 expected_region_highlight=(
   '1 1 builtin' # :
   '3 14 default' # $(( 6 * 9 ))
+  '3 14 arithmetic-expansion' # $(( 6 * 9 ))
 )
-expected_mismatch="currently the actual highlighting has one superfluous group that highlights the asterisk is highlighted as 'globbing'"
diff --git a/highlighters/main/test-data/arith2.zsh b/highlighters/main/test-data/arith2.zsh
index 7e98476..af981d7 100644
--- a/highlighters/main/test-data/arith2.zsh
+++ b/highlighters/main/test-data/arith2.zsh
@@ -34,4 +34,5 @@ expected_region_highlight=(
   '1 1 builtin' # :
   '3 16 default' # "$(( 6 * 9 ))"
   '3 16 double-quoted-argument' # "$(( 6 * 9 ))"
+  '4 15 arithmetic-expansion' # $(( 6 * 9 ))
 )