main: Add arithmetic substitution highlighting

Closes #607 #649 #704

Author: phy1729

highlighters/main/main-highlighter.zsh

@@ -1300,8 +1300,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 substitution
+              (( 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]
@@ -1447,11 +1452,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 substitution
+                  (( 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 ))
@@ -1632,6 +1643,100 @@ _zsh_highlight_main_highlighter_highlight_backtick()
   REPLY=$i
 }
 
+# Highlight special chars inside arithmetic substitutions
+_zsh_highlight_main_highlighter_highlight_arithmetic()
+{
+  local -a match mbegin mend saved_reply
+  local MATCH; integer MBEGIN MEND
+  local i j k ret style
+  integer paren_depth
+  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
+      ')')
+        if (( paren_depth )); then
+          (( paren_depth-- ))
+          continue
+        fi
+        [[ $arg[i+1] == ')' ]] && 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 substitution
+        return 1
+        ;;
+      '(')
+        (( paren_depth++ ))
+        continue
+        ;;
+      [\'\"\\])
+        style=unknown-token
+        ;;
+      '`')
+        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 substitution
+              (( 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,i+1] == '))' ]]; then
+    style=arithmetic-expansion
+    (( i++ ))
+  else
+    # If unclosed, i points past the end
+    (( i-- ))
+    style=arithmetic-expansion-unclosed
+  fi
+  reply=($(( start_pos + $1 - 1)) $(( start_pos + i )) $style $reply)
+  REPLY=$i
+}
+
+
 # Called with a single positional argument.
 # Perform filename expansion (tilde expansion) on the argument and set $REPLY to the expanded value.
 #

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' # )
 )

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 ))
 )