Support indexing with :let

Author: J-Fields

src/cmd_line/commands/let.ts

@@ -8,12 +8,13 @@ import {
   add,
   concat,
   divide,
+  int,
   modulo,
   multiply,
   str,
   subtract,
 } from '../../vimscript/expression/build';
-import { EvaluationContext } from '../../vimscript/expression/evaluate';
+import { EvaluationContext, toInt, toString } from '../../vimscript/expression/evaluate';
 import {
   envVariableParser,
   expressionParser,
@@ -27,7 +28,6 @@ import {
   Expression,
   OptionExpression,
   RegisterExpression,
-  Value,
   VariableExpression,
 } from '../../vimscript/expression/types';
 import { displayValue } from '../../vimscript/expression/displayValue';
@@ -37,6 +37,11 @@ type Unpack = {
   type: 'unpack';
   names: string[];
 };
+type Index = {
+  type: 'index';
+  variable: VariableExpression;
+  index: Expression;
+};
 
 export type LetCommandOperation = '=' | '+=' | '-=' | '*=' | '/=' | '%=' | '.=' | '..=';
 export type LetCommandVariable =
@@ -47,7 +52,7 @@ export type LetCommandVariable =
 export type LetCommandArgs =
   | {
       operation: LetCommandOperation;
-      variable: LetCommandVariable | Unpack;
+      variable: LetCommandVariable | Unpack | Index;
       expression: Expression;
       lock: boolean;
     }
@@ -81,8 +86,16 @@ const unpackParser: Parser<Unpack> = sepBy(varNameParser, string(',').trim(optWh
     names,
   }));
 
+const indexParser: Parser<Index> = seq(
+  variableParser,
+  expressionParser.wrap(string('[').then(optWhitespace), optWhitespace.then(string(']'))),
+).map(([variable, index]) => ({
+  type: 'index',
+  variable,
+  index,
+}));
+
 export class LetCommand extends ExCommand {
-  // TODO: Support indexing
   // TODO: Support slicing
   public static readonly argParser = (lock: boolean) =>
     alt<LetCommand>(
@@ -92,7 +105,7 @@ export class LetCommand extends ExCommand {
       // `:let {var} .= {expr}`
       whitespace.then(
         seq(
-          alt<LetCommandVariable | Unpack>(letVarParser, unpackParser),
+          alt<LetCommandVariable | Unpack | Index>(unpackParser, indexParser, letVarParser),
           operationParser.trim(optWhitespace),
           expressionParser,
         ).map(
@@ -141,30 +154,34 @@ export class LetCommand extends ExCommand {
         }
       }
 
-      let value = context.evaluate(this.args.expression);
-      if (variable.type === 'variable') {
+      const value = context.evaluate(this.args.expression);
+      const newValue = (_var: Expression) => {
         if (this.args.operation === '+=') {
-          value = context.evaluate(add(variable, value));
+          return context.evaluate(add(_var, value));
         } else if (this.args.operation === '-=') {
-          value = context.evaluate(subtract(variable, value));
+          return context.evaluate(subtract(_var, value));
         } else if (this.args.operation === '*=') {
-          value = context.evaluate(multiply(variable, value));
+          return context.evaluate(multiply(_var, value));
         } else if (this.args.operation === '/=') {
-          value = context.evaluate(divide(variable, value));
+          return context.evaluate(divide(_var, value));
         } else if (this.args.operation === '%=') {
-          value = context.evaluate(modulo(variable, value));
+          return context.evaluate(modulo(_var, value));
         } else if (this.args.operation === '.=') {
-          value = context.evaluate(concat(variable, value));
+          return context.evaluate(concat(_var, value));
         } else if (this.args.operation === '..=') {
-          value = context.evaluate(concat(variable, value));
+          return context.evaluate(concat(_var, value));
         }
-        context.setVariable(variable, value, this.args.lock);
+        return value;
+      };
+
+      if (variable.type === 'variable') {
+        context.setVariable(variable, newValue(variable), this.args.lock);
       } else if (variable.type === 'register') {
         // TODO
       } else if (variable.type === 'option') {
         // TODO
       } else if (variable.type === 'env_variable') {
-        value = str(env[variable.name] ?? '');
+        // TODO
       } else if (variable.type === 'unpack') {
         // TODO: Support :let [a, b; rest] = ["aval", "bval", 3, 4]
         if (value.type !== 'list') {
@@ -176,13 +193,33 @@ export class LetCommand extends ExCommand {
         if (variable.names.length > value.items.length) {
           throw VimError.fromCode(ErrorCode.MoreTargetsThanListItems);
         }
-        for (let i = 0; i < variable.names.length; i++) {
-          await new LetCommand({
-            operation: this.args.operation,
-            variable: { type: 'variable', namespace: undefined, name: variable.names[i] },
-            expression: value.items[i],
-            lock: this.args.lock,
-          }).execute(vimState);
+        for (const name of variable.names) {
+          const item: VariableExpression = { type: 'variable', namespace: undefined, name };
+          context.setVariable(item, newValue(item), this.args.lock);
+        }
+      } else if (variable.type === 'index') {
+        const varValue = context.evaluate(variable.variable);
+        if (varValue.type === 'list') {
+          const idx = toInt(context.evaluate(variable.index));
+          const newItem = newValue({
+            type: 'index',
+            expression: variable.variable,
+            index: int(idx),
+          });
+          varValue.items[idx] = newItem;
+          context.setVariable(variable.variable, varValue, this.args.lock);
+        } else if (varValue.type === 'dict_val') {
+          const key = toString(context.evaluate(variable.index));
+          const newItem = newValue({
+            type: 'entry',
+            expression: variable.variable,
+            entryName: key,
+          });
+          varValue.items.set(key, newItem);
+          context.setVariable(variable.variable, varValue, this.args.lock);
+        } else {
+          // TODO: Support blobs
+          throw VimError.fromCode(ErrorCode.CanOnlyIndexAListDictionaryOrBlob);
         }
       }
     }

src/error.ts

@@ -48,6 +48,7 @@ export enum ErrorCode {
   ArgumentOfSortMustBeAList = 686,
   LessTargetsThanListItems = 687,
   MoreTargetsThanListItems = 688,
+  CanOnlyIndexAListDictionaryOrBlob = 689,
   CanOnlyCompareListWithList = 691,
   InvalidOperationForList = 692,
   CannotIndexAFuncref = 695,
@@ -132,6 +133,7 @@ export const ErrorMessage: IErrorMessage = {
   686: 'Argument of sort() must be a List',
   687: 'Less targets than List items',
   688: 'More targets than List items',
+  689: 'Can only index a List, Dictionary or Blob',
   691: 'Can only compare List with List',
   692: 'Invalid operation for List',
   695: 'Cannot index a Funcref',

src/vimscript/expression/evaluate.ts

@@ -23,7 +23,7 @@ import {
 // ID of next lambda; incremented each time one is created
 let lambdaNumber = 1;
 
-function toInt(value: Value): number {
+export function toInt(value: Value): number {
   switch (value.type) {
     case 'number':
       return value.value;

test/vimscript/exCommandParse.test.ts

@@ -379,6 +379,20 @@ suite('Ex command parsing', () => {
       }),
     );
 
+    exParseTest(
+      ':let x[7] = "foo"',
+      new LetCommand({
+        operation: '=',
+        variable: {
+          type: 'index',
+          variable: { type: 'variable', namespace: undefined, name: 'x' },
+          index: int(7),
+        },
+        expression: str('foo'),
+        lock: false,
+      }),
+    );
+
     exParseTest(
       ':const foo = 5',
       new LetCommand({ operation: '=', variable: variable('foo'), expression: int(5), lock: true }),