diff --git a/src/komet/kast/syntax.py b/src/komet/kast/syntax.py index aaba07e..0b1cab3 100644 --- a/src/komet/kast/syntax.py +++ b/src/komet/kast/syntax.py @@ -95,6 +95,10 @@ def sc_bytes(b: bytes) -> KInner: return KApply('SCVal:Bytes', [token(b)]) +def sc_address(address: KInner) -> KInner: + return KApply('SCVal:Address', [address]) + + def sc_vec(l: Iterable[KInner]) -> KInner: return KApply('SCVal:Vec', list_of(l)) diff --git a/src/komet/scval.py b/src/komet/scval.py index 9e3b368..c10de2a 100644 --- a/src/komet/scval.py +++ b/src/komet/scval.py @@ -10,6 +10,9 @@ from pyk.prelude.utils import token from .kast.syntax import ( + account_id, + contract_id, + sc_address, sc_bool, sc_bytes, sc_i32, @@ -120,6 +123,30 @@ def to_kast(self) -> KInner: return sc_bytes(self.val) +@dataclass(frozen=True) +class AccountId: + val: bytes + + def to_kast(self) -> KInner: + return account_id(self.val) + + +@dataclass(frozen=True) +class ContractId: + val: bytes + + def to_kast(self) -> KInner: + return contract_id(self.val) + + +@dataclass(frozen=True) +class SCAddress(SCValue): + val: AccountId | ContractId + + def to_kast(self) -> KInner: + return sc_address(self.val.to_kast()) + + @dataclass(frozen=True) class SCVec(SCValue): val: tuple[SCValue] @@ -150,6 +177,7 @@ def to_kast(self) -> KInner: 'u256': 'SCU256Type', 'symbol': 'SCSymbolType', 'bytes': 'SCBytesType', + 'address': 'SCAddressType', 'vec': 'SCVecType', 'map': 'SCMapType', } @@ -330,6 +358,21 @@ def as_var(cls, name: str) -> tuple[KInner, tuple[KInner, ...]]: return KApply('SCVal:Bytes', [KVariable(name, KSort('Bytes'))]), () +@dataclass +class SCAddressType(SCMonomorphicType): + def strategy(self) -> strategies.SearchStrategy: + def target(p: bool, val: bytes) -> SCAddress: + if p: + return SCAddress(ContractId(val)) + return SCAddress(AccountId(val)) + + return strategies.builds(target, strategies.booleans(), strategies.binary(min_size=32, max_size=32)) + + @classmethod + def as_var(cls, name: str) -> tuple[KInner, tuple[KInner, ...]]: + return KApply('SCVal:Address', [KVariable(name, KSort('Address'))]), () + + @dataclass class SCVecType(SCType): element: SCType