Merge pull request #58 from algorandfoundation/fix/local-state

fix: ensure local state behaviour matches AVM

Author: boblat

examples/local-storage/contract.algo.ts

@@ -0,0 +1,163 @@
+import type { Account, bytes, uint64 } from '@algorandfoundation/algorand-typescript'
+import { Bytes, Global, LocalState, Txn, arc4, assert, contract } from '@algorandfoundation/algorand-typescript'
+
+/**
+ * A contract demonstrating local storage functionality.
+ * This contract shows how to use local state storage in an Algorand smart contract,
+ * including initialization, reading, writing, and clearing of local state values.
+ * Local state is per-account storage that requires accounts to opt-in before use.
+ *
+ * @stateTotals.localBytes - 4 bytes allocated for local byte storage
+ * @stateTotals.localUints - 3 uints allocated for local integer storage
+ */
+@contract({ stateTotals: { localBytes: 4, localUints: 3 } })
+export default class LocalStorage extends arc4.Contract {
+  // example: INIT_LOCAL_STATE
+  public localInt = LocalState<uint64>({ key: 'int' })
+  public localIntNoDefault = LocalState<uint64>()
+  public localBytes = LocalState<bytes>()
+  public localString = LocalState<string>()
+  public localBool = LocalState<boolean>()
+  public localAccount = LocalState<Account>()
+  // example: INIT_LOCAL_STATE
+
+  // example: OPT_IN_TO_APPLICATION
+  /**
+   * Initializes local state values when an account opts into the application.
+   * This method can only be called during an OptIn transaction.
+   * Sets initial values for all local state variables:
+   * - localInt: 100
+   * - localIntNoDefault: 200
+   * - localBytes: 'Silvio'
+   * - localString: 'Micali'
+   * - localBool: true
+   * - localAccount: sender's address
+   */
+  @arc4.abimethod({ allowActions: 'OptIn' })
+  public optInToApplication(): void {
+    this.localInt(Txn.sender).value = 100
+    this.localIntNoDefault(Txn.sender).value = 200
+    this.localBytes(Txn.sender).value = Bytes('Silvio')
+    this.localString(Txn.sender).value = 'Micali'
+    this.localBool(Txn.sender).value = true
+    this.localAccount(Txn.sender).value = Txn.sender
+  }
+  // example: OPT_IN_TO_APPLICATION
+
+  // example: READ_LOCAL_STATE
+  /**
+   * Reads and returns all local state values for the transaction sender.
+   * @returns A tuple containing:
+   * - [0] uint64: The value of localInt
+   * - [1] uint64: The value of localIntNoDefault
+   * - [2] bytes: The value of localBytes
+   * - [3] string: The value of localString
+   * - [4] boolean: The value of localBool
+   * - [5] Address: The value of localAccount converted to Address type
+   */
+  @arc4.abimethod({ readonly: true })
+  public readLocalState(): [uint64, uint64, bytes, string, boolean, arc4.Address] {
+    const sender = Txn.sender
+    // Convert Account reference type to native Address type for return value
+    const accountAddress = new arc4.Address(this.localAccount(sender).value)
+
+    return [
+      this.localInt(sender).value,
+      this.localIntNoDefault(sender).value,
+      this.localBytes(sender).value,
+      this.localString(sender).value,
+      this.localBool(sender).value,
+      accountAddress,
+    ]
+  }
+  // example: READ_LOCAL_STATE
+
+  // example: WRITE_LOCAL_STATE
+  /**
+   * Updates multiple local state values for the transaction sender.
+   * Requires the account to be opted into the application.
+   * @param valueString - New string value to store
+   * @param valueBool - New boolean value to store
+   * @param valueAccount - New account address to store
+   */
+  @arc4.abimethod()
+  public writeLocalState(valueString: string, valueBool: boolean, valueAccount: Account): void {
+    // Dynamic keys must be explicitly reserved in the contract's stateTotals configuration
+    const sender = Txn.sender
+
+    assert(sender.isOptedIn(Global.currentApplicationId), 'Account must opt in to contract first')
+
+    this.localString(sender).value = valueString
+    this.localBool(sender).value = valueBool
+    this.localAccount(sender).value = valueAccount
+
+    assert(this.localString(sender).value === valueString)
+    assert(this.localBool(sender).value === valueBool)
+    assert(this.localAccount(sender).value === valueAccount)
+  }
+  // example: WRITE_LOCAL_STATE
+
+  // example: WRITE_DYNAMIC_LOCAL_STATE
+  /**
+   * Writes a value to local state using a dynamic key.
+   * Demonstrates dynamic key-value storage in local state.
+   * @param key - The dynamic key to store the value under
+   * @param value - The string value to store
+   * @returns The stored string value
+   */
+  @arc4.abimethod()
+  public writeDynamicLocalState(key: string, value: string): string {
+    const sender = Txn.sender
+
+    assert(sender.isOptedIn(Global.currentApplicationId), 'Account must opt in to contract first')
+
+    const localDynamicAccess = LocalState<string>({ key })
+
+    localDynamicAccess(sender).value = value
+
+    assert(localDynamicAccess(sender).value === value)
+
+    return localDynamicAccess(sender).value
+  }
+  // example: WRITE_DYNAMIC_LOCAL_STATE
+
+  // example: READ_DYNAMIC_LOCAL_STATE
+  /**
+   * Reads a value from local state using a dynamic key.
+   * @param key - The dynamic key to read the value from
+   * @returns The stored string value for the given key
+   */
+  @arc4.abimethod()
+  public readDynamicLocalState(key: string): string {
+    const sender = Txn.sender
+
+    assert(sender.isOptedIn(Global.currentApplicationId), 'Account must opt in to contract first')
+
+    const localDynamicAccess = LocalState<string>({ key })
+
+    assert(localDynamicAccess(sender).hasValue, 'Key not found')
+
+    return localDynamicAccess(sender).value
+  }
+  // example: READ_DYNAMIC_LOCAL_STATE
+
+  // example: CLEAR_LOCAL_STATE
+  /**
+   * Clears all local state values for the transaction sender.
+   * After calling this method, all local state values will be deleted.
+   */
+  @arc4.abimethod()
+  public clearLocalState(): void {
+    const sender = Txn.sender
+
+    assert(sender.isOptedIn(Global.currentApplicationId), 'Account must opt in to contract first')
+
+    this.localInt(sender).delete()
+    this.localIntNoDefault(sender).delete()
+    this.localBytes(sender).delete()
+    this.localString(sender).delete()
+    this.localBool(sender).delete()
+    this.localAccount(sender).delete()
+  }
+  // example: CLEAR_LOCAL_STATE
+}

examples/local-storage/contract.spec.ts

@@ -0,0 +1,95 @@
+import { Bytes } from '@algorandfoundation/algorand-typescript'
+import { TestExecutionContext } from '@algorandfoundation/algorand-typescript-testing'
+import { describe, expect, it } from 'vitest'
+import LocalStorage from './contract.algo'
+
+describe('LocalStorage contract', () => {
+  const ctx = new TestExecutionContext()
+
+  it('should initialize local state with correct values after opting in', () => {
+    const contract = ctx.contract.create(LocalStorage)
+
+    contract.optInToApplication()
+
+    const [localInt, localIntNoDefault, localBytes, localString, localBool, localAddress] = contract.readLocalState()
+
+    expect(localInt.valueOf()).toBe(100)
+    expect(localIntNoDefault.valueOf()).toBe(200)
+    expect(localBytes.toString()).toBe(Bytes('Silvio').toString())
+    expect(localString).toBe('Micali')
+    expect(localBool).toBe(true)
+    expect(localAddress.bytes.toString()).toBe(ctx.defaultSender.bytes.toString())
+  })
+
+  it('should write and verify multiple local state values', () => {
+    const contract = ctx.contract.create(LocalStorage)
+    const account = ctx.any.account({ optedApplications: [ctx.ledger.getApplicationForContract(contract)] })
+
+    ctx.txn.createScope([ctx.any.txn.applicationCall({ appId: contract, sender: account, onCompletion: 'OptIn' })]).execute(() => {
+      contract.optInToApplication()
+    })
+
+    ctx.txn.createScope([ctx.any.txn.applicationCall({ appId: contract, sender: account })]).execute(() => {
+      contract.writeLocalState('New String', false, account)
+
+      const [, , , localString, localBool, localAccount] = contract.readLocalState()
+
+      expect(localString).toBe('New String')
+      expect(localBool).toBe(false)
+      expect(localAccount.bytes).toEqual(account.bytes)
+    })
+  })
+
+  it('should write and read dynamic local state values', () => {
+    const contract = ctx.contract.create(LocalStorage)
+    const account = ctx.any.account({ optedApplications: [ctx.ledger.getApplicationForContract(contract)] })
+
+    ctx.txn.createScope([ctx.any.txn.applicationCall({ appId: contract, sender: account, onCompletion: 'OptIn' })]).execute(() => {
+      contract.optInToApplication()
+    })
+
+    ctx.txn.createScope([ctx.any.txn.applicationCall({ appId: contract, sender: account })]).execute(() => {
+      const testKey = 'testKey'
+      const testValue = 'testValue'
+
+      const writtenValue = contract.writeDynamicLocalState(testKey, testValue)
+      expect(writtenValue).toBe(testValue)
+
+      const readValue = contract.readDynamicLocalState(testKey)
+      expect(readValue).toBe(testValue)
+    })
+  })
+
+  it('should clear all local state values', () => {
+    const contract = ctx.contract.create(LocalStorage)
+    const account = ctx.any.account({ optedApplications: [ctx.ledger.getApplicationForContract(contract)] })
+    ctx.txn.createScope([ctx.any.txn.applicationCall({ appId: contract, sender: account, onCompletion: 'OptIn' })]).execute(() => {
+      contract.optInToApplication()
+    })
+
+    ctx.txn.createScope([ctx.any.txn.applicationCall({ appId: contract, sender: account })]).execute(() => {
+      contract.clearLocalState()
+
+      expect(() => contract.readLocalState()).toThrow()
+    })
+  })
+
+  it('should fail to write local state if not opted in', () => {
+    const contract = ctx.contract.create(LocalStorage)
+    const newAccount = ctx.defaultSender
+
+    expect(() => contract.writeLocalState('New String', false, newAccount)).toThrow('Account must opt in to contract first')
+  })
+
+  it('should fail to read dynamic local state for non-existent key', () => {
+    const contract = ctx.contract.create(LocalStorage)
+    const account = ctx.any.account({ optedApplications: [ctx.ledger.getApplicationForContract(contract)] })
+    ctx.txn.createScope([ctx.any.txn.applicationCall({ appId: contract, sender: account, onCompletion: 'OptIn' })]).execute(() => {
+      contract.optInToApplication()
+    })
+
+    ctx.txn.createScope([ctx.any.txn.applicationCall({ appId: contract, sender: account })]).execute(() => {
+      expect(() => contract.readDynamicLocalState('nonexistentKey')).toThrow('Key not found')
+    })
+  })
+})

package-lock.json

@@ -71,8 +71,8 @@
     "tslib": "^2.6.2"
   },
   "dependencies": {
-    "@algorandfoundation/algorand-typescript": "^1.0.0-beta.16",
-    "@algorandfoundation/puya-ts": "^1.0.0-beta.23",
+    "@algorandfoundation/algorand-typescript": "^1.0.0-beta.17",
+    "@algorandfoundation/puya-ts": "^1.0.0-beta.24",
     "elliptic": "^6.5.7",
     "js-sha256": "^0.11.0",
     "js-sha3": "^0.9.3",

package.json

@@ -8,6 +8,7 @@ import type {
 } from '@algorandfoundation/algorand-typescript'
 import { encodingUtil } from '@algorandfoundation/puya-ts'
 import js_sha512 from 'js-sha512'
+import type { AccountMap } from '../collections/custom-key-map'
 import { BytesMap, Uint64Map } from '../collections/custom-key-map'
 import {
   ALGORAND_ADDRESS_BYTE_LENGTH,
@@ -26,7 +27,7 @@ import { asBigInt, asBytes, asUint64, asUint64Cls, asUint8Array, conactUint8Arra
 import { BytesBackedCls, Uint64BackedCls } from './base'
 import type { StubUint64Compat } from './primitives'
 import { Bytes, BytesCls, Uint64Cls } from './primitives'
-import type { GlobalStateCls } from './state'
+import type { GlobalStateCls, LocalStateCls } from './state'
 
 export class AssetHolding {
   balance: uint64
@@ -133,6 +134,7 @@ export class ApplicationData {
     appLogs: bytes[]
     globalStates: BytesMap<GlobalStateCls<unknown>>
     localStates: BytesMap<LocalState<unknown>>
+    localStateMaps: BytesMap<AccountMap<LocalStateCls<unknown>>>
     boxes: BytesMap<Uint8Array>
   }
 
@@ -151,6 +153,7 @@ export class ApplicationData {
       appLogs: [],
       globalStates: new BytesMap(),
       localStates: new BytesMap(),
+      localStateMaps: new BytesMap(),
       boxes: new BytesMap(),
     }
   }

src/impl/reference.ts

@@ -91,13 +91,27 @@ export class LocalStateCls<ValueType> {
 }
 
 export class LocalStateMapCls<ValueType> {
-  #value = new AccountMap<LocalStateCls<ValueType>>()
+  private applicationId: uint64
 
-  getValue(account: Account): LocalStateCls<ValueType> {
-    if (!this.#value.has(account)) {
-      this.#value.set(account, new LocalStateCls<ValueType>())
+  constructor() {
+    this.applicationId = lazyContext.activeGroup.activeApplicationId
+  }
+
+  getValue(key: string | bytes | undefined, account: Account): LocalStateCls<ValueType> {
+    const bytesKey = key === undefined ? Bytes() : asBytes(key)
+    const localStateMap = this.ensureApplicationLocalStateMap(bytesKey)
+    if (!localStateMap.has(account)) {
+      localStateMap.set(account, new LocalStateCls())
+    }
+    return localStateMap.getOrFail(account) as LocalStateCls<ValueType>
+  }
+
+  private ensureApplicationLocalStateMap(key: bytes | string) {
+    const applicationData = lazyContext.ledger.applicationDataMap.getOrFail(this.applicationId)!.application
+    if (!applicationData.localStateMaps.has(key)) {
+      applicationData.localStateMaps.set(key, new AccountMap<LocalStateCls<ValueType>>())
     }
-    return this.#value.getOrFail(account)!
+    return applicationData.localStateMaps.getOrFail(key)
   }
 }
 
@@ -107,7 +121,7 @@ export function GlobalState<ValueType>(options?: GlobalStateOptions<ValueType>):
 
 export function LocalState<ValueType>(options?: { key?: bytes | string }): LocalStateType<ValueType> {
   function localStateInternal(account: Account): LocalStateForAccount<ValueType> {
-    return localStateInternal.map.getValue(account)
+    return localStateInternal.map.getValue(localStateInternal.key, account)
   }
   localStateInternal.key = options?.key
   localStateInternal.hasKey = options?.key !== undefined && options.key.length > 0