feat: Array push, pop, slice, and at

src/awst_build/eb/arc4/arrays.ts

@@ -2,7 +2,6 @@ import { nodeFactory } from '../../../awst/node-factory'
 import type { Expression } from '../../../awst/nodes'
 import type { SourceLocation } from '../../../awst/source-location'
 import { uint64WType } from '../../../awst/wtypes'
-import { logger } from '../../../logger'
 import { codeInvariant, invariant } from '../../../util'
 import type { PType } from '../../ptypes'
 import { IterableIteratorType, NumericLiteralPType, TuplePType, uint64PType } from '../../ptypes'
@@ -15,11 +14,12 @@ import {
 } from '../../ptypes/arc4-types'
 import { instanceEb } from '../../type-registry'
 import type { InstanceBuilder } from '../index'
-import { BuilderBinaryOp, FunctionBuilder, InstanceExpressionBuilder, NodeBuilder } from '../index'
+import { FunctionBuilder, InstanceExpressionBuilder, NodeBuilder } from '../index'
 import { IterableIteratorExpressionBuilder } from '../iterable-iterator-expression-builder'
-import { BigIntLiteralExpressionBuilder } from '../literal/big-int-literal-expression-builder'
+import { AtFunctionBuilder } from '../shared/at-function-builder'
+import { SliceFunctionBuilder } from '../shared/slice-function-builder'
 import { UInt64ExpressionBuilder } from '../uint64-expression-builder'
-import { requireExpressionOfType, requireInstanceBuilder } from '../util'
+import { requireExpressionOfType } from '../util'
 import { parseFunctionArgs } from '../util/arg-parsing'
 
 export class DynamicArrayConstructorBuilder extends NodeBuilder {
@@ -101,11 +101,19 @@ export abstract class ArrayExpressionBuilder<
   memberAccess(name: string, sourceLocation: SourceLocation): NodeBuilder {
     switch (name) {
       case 'at':
-        return new ArrayAtFunctionBuilder(this)
+        return new AtFunctionBuilder(
+          this.resolve(),
+          this.ptype.elementType,
+          this.ptype instanceof StaticArrayType
+            ? this.ptype.arraySize
+            : requireExpressionOfType(this.memberAccess('length', sourceLocation), uint64PType),
+        )
       case 'entries':
         return new EntriesFunctionBuilder(this)
       case 'copy':
         return new CopyFunctionBuilder(this)
+      case 'slice':
+        return new SliceFunctionBuilder(this.resolve(), this.ptype)
     }
     return super.memberAccess(name, sourceLocation)
   }
@@ -117,6 +125,7 @@ class CopyFunctionBuilder extends FunctionBuilder {
   }
 
   call(args: ReadonlyArray<InstanceBuilder>, typeArgs: ReadonlyArray<PType>, sourceLocation: SourceLocation): NodeBuilder {
+    parseFunctionArgs({ args, typeArgs, genericTypeArgs: 0, argSpec: (a) => [], funcName: 'copy', callLocation: sourceLocation })
     return instanceEb(
       nodeFactory.copy({
         value: this.arrayBuilder.resolve(),
@@ -164,6 +173,10 @@ export class DynamicArrayExpressionBuilder extends ArrayExpressionBuilder<Dynami
             wtype: uint64WType,
           }),
         )
+      case 'push':
+        return new ArrayPushFunctionBuilder(this)
+      case 'pop':
+        return new ArrayPopFunctionBuilder(this)
     }
     return super.memberAccess(name, sourceLocation)
   }
@@ -184,58 +197,61 @@ export class StaticArrayExpressionBuilder extends ArrayExpressionBuilder<StaticA
   }
 }
 
-export class ArrayAtFunctionBuilder extends FunctionBuilder {
-  constructor(private arrayBuilder: ArrayExpressionBuilder<StaticArrayType | DynamicArrayType>) {
+export class ArrayPushFunctionBuilder extends FunctionBuilder {
+  constructor(private arrayBuilder: ArrayExpressionBuilder<DynamicArrayType>) {
     super(arrayBuilder.sourceLocation)
   }
 
   call(args: ReadonlyArray<InstanceBuilder>, typeArgs: ReadonlyArray<PType>, sourceLocation: SourceLocation): NodeBuilder {
+    const elementType = this.arrayBuilder.ptype.elementType
     const {
-      args: [index],
+      args: [...items],
     } = parseFunctionArgs({
       args,
       typeArgs,
       funcName: 'at',
       callLocation: sourceLocation,
       genericTypeArgs: 0,
-      argSpec: (a) => [a.required()],
+      argSpec: (a) => [a.required(elementType), ...args.slice(1).map(() => a.required(elementType))],
+    })
+
+    return instanceEb(
+      nodeFactory.arrayExtend({
+        base: this.arrayBuilder.resolve(),
+        other: nodeFactory.tupleExpression({
+          items: items.map((i) => i.resolve()),
+          sourceLocation,
+        }),
+        sourceLocation,
+        wtype: this.arrayBuilder.ptype.wtype,
+      }),
+      this.arrayBuilder.ptype,
+    )
+  }
+}
+export class ArrayPopFunctionBuilder extends FunctionBuilder {
+  constructor(private arrayBuilder: ArrayExpressionBuilder<DynamicArrayType>) {
+    super(arrayBuilder.sourceLocation)
+  }
+
+  call(args: ReadonlyArray<InstanceBuilder>, typeArgs: ReadonlyArray<PType>, sourceLocation: SourceLocation): NodeBuilder {
+    const elementType = this.arrayBuilder.ptype.elementType
+    parseFunctionArgs({
+      args,
+      typeArgs,
+      funcName: 'at',
+      callLocation: sourceLocation,
+      genericTypeArgs: 0,
+      argSpec: () => [],
     })
 
-    let indexExpr: Expression
-
-    if (index instanceof BigIntLiteralExpressionBuilder) {
-      if (this.arrayBuilder.ptype instanceof StaticArrayType) {
-        const staticSize = this.arrayBuilder.ptype.arraySize
-        if (index.value < -staticSize || index.value >= staticSize) {
-          logger.error(index.sourceLocation, 'Index access outside the bounds of the array')
-        }
-        indexExpr = nodeFactory.uInt64Constant({
-          value: index.value < 0 ? staticSize + index.value : index.value,
-          sourceLocation: index.sourceLocation,
-        })
-      } else {
-        const dynamicSize = requireInstanceBuilder(this.arrayBuilder.memberAccess('length', sourceLocation))
-        const absoluteIndex = nodeFactory.uInt64Constant({
-          value: index.value < 0n ? index.value * -1n : index.value,
-          sourceLocation: index.sourceLocation,
-        })
-        if (index.value < 0) {
-          indexExpr = dynamicSize.binaryOp(new UInt64ExpressionBuilder(absoluteIndex), BuilderBinaryOp.sub, index.sourceLocation).resolve()
-        } else {
-          indexExpr = absoluteIndex
-        }
-      }
-    } else {
-      indexExpr = index.resolveToPType(uint64PType).resolve()
-    }
     return instanceEb(
-      nodeFactory.indexExpression({
+      nodeFactory.arrayPop({
         base: this.arrayBuilder.resolve(),
-        index: indexExpr,
-        wtype: this.arrayBuilder.ptype.elementType.wtype,
         sourceLocation,
+        wtype: elementType.wtype,
       }),
-      this.arrayBuilder.ptype.elementType,
+      elementType,
     )
   }
 }

src/awst_build/eb/bytes-expression-builder.ts

@@ -3,20 +3,12 @@ import { wtypes } from '../../awst'
 import { intrinsicFactory } from '../../awst/intrinsic-factory'
 import { nodeFactory } from '../../awst/node-factory'
 import type { Expression } from '../../awst/nodes'
-import {
-  BytesBinaryOperator,
-  BytesConstant,
-  BytesEncoding,
-  BytesUnaryOperator,
-  IntegerConstant,
-  StringConstant,
-  UInt64BinaryOperator,
-} from '../../awst/nodes'
+import { BytesBinaryOperator, BytesConstant, BytesEncoding, BytesUnaryOperator, IntegerConstant, StringConstant } from '../../awst/nodes'
 import type { SourceLocation } from '../../awst/source-location'
 import { stringWType } from '../../awst/wtypes'
 import { CodeError, wrapInCodeError } from '../../errors'
 import { logger } from '../../logger'
-import { base32ToUint8Array, base64ToUint8Array, hexToUint8Array, invariant, uint8ArrayToUtf8, utf8ToUint8Array } from '../../util'
+import { base32ToUint8Array, base64ToUint8Array, hexToUint8Array, uint8ArrayToUtf8, utf8ToUint8Array } from '../../util'
 import type { InstanceType, PType } from '../ptypes'
 import {
   ArrayPType,
@@ -29,11 +21,12 @@ import {
   stringPType,
   uint64PType,
 } from '../ptypes'
-import { instanceEb } from '../type-registry'
 import type { BuilderComparisonOp, InstanceBuilder, NodeBuilder } from './index'
 import { BuilderUnaryOp, FunctionBuilder, InstanceExpressionBuilder, ParameterlessFunctionBuilder } from './index'
 import { ArrayLiteralExpressionBuilder } from './literal/array-literal-expression-builder'
 import { BigIntLiteralExpressionBuilder } from './literal/big-int-literal-expression-builder'
+import { AtFunctionBuilder } from './shared/at-function-builder'
+import { SliceFunctionBuilder } from './shared/slice-function-builder'
 import { StringExpressionBuilder } from './string-expression-builder'
 import { UInt64ExpressionBuilder } from './uint64-expression-builder'
 import { requireExpressionOfType, requireExpressionsOfType } from './util'
@@ -203,7 +196,13 @@ export class BytesExpressionBuilder extends InstanceExpressionBuilder<InstanceTy
       case 'concat':
         return new ConcatExpressionBuilder(this._expr)
       case 'at':
-        return new BytesAtBuilder(this._expr)
+        return new AtFunctionBuilder(
+          this._expr,
+          bytesPType,
+          requireExpressionOfType(this.memberAccess('length', sourceLocation), uint64PType),
+        )
+      case 'slice':
+        return new SliceFunctionBuilder(this._expr, bytesPType)
     }
     return super.memberAccess(name, sourceLocation)
   }
@@ -272,44 +271,6 @@ export class BytesInvertBuilder extends ParameterlessFunctionBuilder {
   }
 }
 
-export class BytesSliceBuilder extends FunctionBuilder {
-  constructor(private expr: awst.Expression) {
-    super(expr.sourceLocation)
-  }
-  call(args: ReadonlyArray<InstanceBuilder>, typeArgs: ReadonlyArray<PType>, sourceLocation: SourceLocation) {
-    const {
-      args: [start, stop],
-    } = parseFunctionArgs({
-      args,
-      typeArgs,
-      genericTypeArgs: 0,
-      callLocation: sourceLocation,
-      funcName: 'slice',
-      argSpec: (a) => [a.optional(uint64PType, numberPType), a.optional(uint64PType, numberPType)],
-    })
-
-    return new BytesExpressionBuilder(
-      nodeFactory.intersectionSliceExpression({
-        base: this.expr,
-        sourceLocation: sourceLocation,
-        beginIndex: start ? getBigIntOrUint64Expr(start) : null,
-        endIndex: stop ? getBigIntOrUint64Expr(stop) : null,
-        wtype: wtypes.bytesWType,
-      }),
-    )
-  }
-}
-
-function getBigIntOrUint64Expr(builder: InstanceBuilder) {
-  if (builder.ptype.equals(numberPType)) {
-    invariant(builder instanceof BigIntLiteralExpressionBuilder, 'Builder for number type must be BigIntLiteral')
-    return builder.value
-  } else {
-    invariant(builder.ptype.equals(uint64PType), 'Builder must be uint64 if not number')
-    return builder.resolve()
-  }
-}
-
 export class ToStringBuilder extends ParameterlessFunctionBuilder {
   constructor(private expr: awst.Expression) {
     super(
@@ -325,57 +286,3 @@ export class ToStringBuilder extends ParameterlessFunctionBuilder {
     )
   }
 }
-
-export class BytesAtBuilder extends FunctionBuilder {
-  constructor(private expr: awst.Expression) {
-    super(expr.sourceLocation)
-  }
-
-  call(args: ReadonlyArray<InstanceBuilder>, typeArgs: ReadonlyArray<PType>, sourceLocation: SourceLocation) {
-    const {
-      args: [index],
-    } = parseFunctionArgs({
-      args,
-      typeArgs,
-      genericTypeArgs: 0,
-      callLocation: sourceLocation,
-      funcName: 'at',
-      argSpec: (a) => [a.required(uint64PType, numberPType)],
-    })
-
-    let indexExpr: Expression
-
-    if (index.ptype.equals(numberPType)) {
-      invariant(index instanceof BigIntLiteralExpressionBuilder, 'Builder for number type must be BigIntLiteral')
-
-      if (index.value < 0) {
-        indexExpr = nodeFactory.uInt64BinaryOperation({
-          op: UInt64BinaryOperator.sub,
-          left: intrinsicFactory.bytesLen({
-            value: this.expr,
-            sourceLocation,
-          }),
-          right: nodeFactory.uInt64Constant({
-            value: index.value * -1n,
-            sourceLocation,
-          }),
-          sourceLocation,
-        })
-      } else {
-        indexExpr = index.resolveToPType(uint64PType).resolve()
-      }
-    } else {
-      indexExpr = index.resolve()
-    }
-
-    return instanceEb(
-      nodeFactory.indexExpression({
-        base: this.expr,
-        sourceLocation: sourceLocation,
-        index: indexExpr,
-        wtype: wtypes.bytesWType,
-      }),
-      bytesPType,
-    )
-  }
-}

src/awst_build/eb/shared/at-function-builder.ts

@@ -0,0 +1,74 @@
+import { nodeFactory } from '../../../awst/node-factory'
+import type { Expression } from '../../../awst/nodes'
+import { UInt64BinaryOperator } from '../../../awst/nodes'
+import type { SourceLocation } from '../../../awst/source-location'
+import type { WType } from '../../../awst/wtypes'
+import { logger } from '../../../logger'
+import type { PType } from '../../ptypes'
+import { numberPType, uint64PType } from '../../ptypes'
+import { instanceEb } from '../../type-registry'
+import { FunctionBuilder, type InstanceBuilder } from '../index'
+import { getBigIntOrUint64Expr } from '../util'
+import { parseFunctionArgs } from '../util/arg-parsing'
+
+export class AtFunctionBuilder extends FunctionBuilder {
+  constructor(
+    private expr: Expression,
+    private itemPType: PType & { wtype: WType },
+    private exprLength: Expression | bigint,
+  ) {
+    super(expr.sourceLocation)
+  }
+
+  call(args: ReadonlyArray<InstanceBuilder>, typeArgs: ReadonlyArray<PType>, sourceLocation: SourceLocation) {
+    const {
+      args: [index],
+    } = parseFunctionArgs({
+      args,
+      typeArgs,
+      genericTypeArgs: 0,
+      callLocation: sourceLocation,
+      funcName: 'at',
+      argSpec: (a) => [a.required(uint64PType, numberPType)],
+    })
+
+    const indexParam = getBigIntOrUint64Expr(index)
+    let indexExpr: Expression
+
+    if (typeof indexParam === 'bigint') {
+      if (typeof this.exprLength === 'bigint') {
+        let indexValue = indexParam < 0 ? this.exprLength + indexParam : indexParam
+        if (indexValue < 0n || indexValue >= this.exprLength) {
+          logger.warn(index.sourceLocation, 'Index access out of bounds')
+          indexValue = 0n
+        }
+        indexExpr = nodeFactory.uInt64Constant({
+          value: indexValue,
+          sourceLocation: index.sourceLocation,
+        })
+      } else {
+        indexExpr = nodeFactory.uInt64BinaryOperation({
+          op: UInt64BinaryOperator.sub,
+          left: this.exprLength,
+          right: nodeFactory.uInt64Constant({
+            value: indexParam * -1n,
+            sourceLocation: index.sourceLocation,
+          }),
+          sourceLocation: index.sourceLocation,
+        })
+      }
+    } else {
+      indexExpr = indexParam
+    }
+
+    return instanceEb(
+      nodeFactory.indexExpression({
+        base: this.expr,
+        sourceLocation: sourceLocation,
+        index: indexExpr,
+        wtype: this.itemPType.wtype,
+      }),
+      this.itemPType,
+    )
+  }
+}