diff --git a/wake/lsp/methods.py b/wake/lsp/methods.py index ef3f0797..5e310483 100644 --- a/wake/lsp/methods.py +++ b/wake/lsp/methods.py @@ -126,6 +126,8 @@ class RequestMethodEnum(StrEnum): # Sake SAKE_PING = "wake/sake/ping" + SAKE_LOAD_WORKSPACE_STATE = "wake/sake/loadWorkspaceState" + SAKE_SAVE_WORKSPACE_STATE = "wake/sake/saveWorkspaceState" SAKE_CREATE_CHAIN = "wake/sake/createChain" SAKE_CONNECT_CHAIN = "wake/sake/connectChain" SAKE_DISCONNECT_CHAIN = "wake/sake/disconnectChain" diff --git a/wake/lsp/sake.py b/wake/lsp/sake.py index 02fb1b89..1367cc3c 100644 --- a/wake/lsp/sake.py +++ b/wake/lsp/sake.py @@ -15,14 +15,14 @@ from wake.config import WakeConfig from wake.development.call_trace import CallTrace from wake.development.chain_interfaces import AnvilChainInterface -from wake.development.core import RequestType, get_fqn_from_address +from wake.development.core import Library, RequestType, get_fqn_from_address from wake.development.globals import set_config from wake.development.json_rpc import JsonRpcError from wake.development.transactions import TransactionStatusEnum from wake.development.utils import ( + AbiNotFound, get_logic_contract, get_name_abi_from_explorer, - AbiNotFound, ) from wake.lsp.context import LspContext from wake.lsp.exceptions import LspError @@ -57,11 +57,6 @@ class SakeCompilationResult(SakeResult): skipped: Dict[str, SkippedInfo] -class ContractInfo(NamedTuple): - abi: List - bytecode: str - - # @dev used for api to include name in json, otherwise tuple is converted to array class ContractInfoLsp(LspModel): abi: List @@ -119,6 +114,7 @@ class SakeCallResult(SakeResult): class SakeTransactParams(SakeCallParams): pass + class SakeTransactResult(SakeResult): return_value: Optional[str] # raw hex encoded bytes, None for Halt error: Optional[str] # user-friendly error string, None for Success @@ -174,8 +170,23 @@ class SakeGetAbiResult(SakeResult): implementation_address: Optional[str] +class WorkspaceState(LspModel): + info_by_fqn: Dict[str, ContractInfo] + libraries: Dict[str, str] + creation_code_index: List[Tuple[Tuple[Tuple[int, str], ...], str]] + fqn_by_metadata: Dict[str, str] + + +class SakeLoadWorkspaceStateParams(SakeParams): + state: WorkspaceState + + +class SakeSaveWorkspaceStateResult(SakeResult): + state: WorkspaceState + + @dataclass -class ContractAbiName: +class ContractInfo: name: str abi: Dict[Union[bytes, Literal["constructor", "fallback", "receive"]], List] @@ -200,8 +211,8 @@ async def wrapper(context: SakeContext, params: SakeParams, *args, **kwargs): class SakeContext: lsp_context: LspContext chains: Dict[str, Tuple[Chain, ContextManager]] - compilation: Dict[str, ContractInfo] - info_by_fqn: Dict[str, ContractAbiName] + bytecode_by_fqn: Dict[str, str] + info_by_fqn: Dict[str, ContractInfo] libraries: Dict[bytes, str] # lib_id -> fqn creation_code_index: List[Tuple[Tuple[Tuple[int, bytes], ...], str]] fqn_by_metadata: Dict[bytes, str] @@ -209,15 +220,58 @@ class SakeContext: def __init__(self, lsp_context: LspContext): self.lsp_context = lsp_context self.chains = {} - self.compilation = {} - self.info_by_fqn = {} - self.libraries = {} - self.creation_code_index = [] - self.fqn_by_metadata = {} + self.bytecode_by_fqn = {} # NOT saved in dump + self.info_by_fqn = {} # saved in dump + self.libraries = {} # saved in dump + self.creation_code_index = [] # saved in dump + self.fqn_by_metadata = {} # saved in dump async def ping(self) -> SakeResult: return SakeResult(success=True) + async def load_workspace_state( + self, params: SakeLoadWorkspaceStateParams + ) -> SakeResult: + try: + self.info_by_fqn = params.state.info_by_fqn + self.libraries = { + bytes.fromhex(lib_id): fqn + for lib_id, fqn in params.state.libraries.items() + } + self.creation_code_index = [ + (tuple((length, bytes.fromhex(h)) for length, h in segments), fqn) + for segments, fqn in params.state.creation_code_index + ] + self.fqn_by_metadata = { + bytes.fromhex(metadata): fqn + for metadata, fqn in params.state.fqn_by_metadata.items() + } + return SakeResult(success=True) + except Exception as e: + raise LspError(ErrorCodes.InternalError, str(e)) from None + + async def save_workspace_state(self) -> SakeSaveWorkspaceStateResult: + try: + return SakeSaveWorkspaceStateResult( + success=True, + state=WorkspaceState( + info_by_fqn=self.info_by_fqn, + libraries={ + lib_id.hex(): fqn for lib_id, fqn in self.libraries.items() + }, + creation_code_index=[ + (tuple((length, h.hex()) for length, h in segments), fqn) + for segments, fqn in self.creation_code_index + ], + fqn_by_metadata={ + metadata.hex(): fqn + for metadata, fqn in self.fqn_by_metadata.items() + }, + ), + ) + except Exception as e: + raise LspError(ErrorCodes.InternalError, str(e)) from None + async def create_chain( self, params: SakeCreateChainParams ) -> SakeCreateChainResult: @@ -356,7 +410,7 @@ async def compile(self) -> SakeCompilationResult: ) = await self.lsp_context.compiler.bytecode_compile() self.libraries.clear() - _compilation = {} + bytecode_by_fqn = {} for fqn, info in contract_info.items(): if ( @@ -421,7 +475,7 @@ async def compile(self) -> SakeCompilationResult: else: raise ValueError(f"Unknown ABI item type: {item['type']}") - self.info_by_fqn[fqn] = ContractAbiName( + self.info_by_fqn[fqn] = ContractInfo( name=contract_node["name"], abi=abi_by_selector ) @@ -449,24 +503,22 @@ async def compile(self) -> SakeCompilationResult: bytes.fromhex(info.evm.deployed_bytecode.object[-106:]) ] = fqn - _compilation[fqn] = ContractInfo( - abi=info.abi, - bytecode=bytecode, - ) + bytecode_by_fqn[fqn] = bytecode wake.development.core.creation_code_index = self.creation_code_index wake.development.core.contracts_by_metadata = self.fqn_by_metadata if success: - self.compilation = _compilation + self.bytecode_by_fqn = bytecode_by_fqn return SakeCompilationResult( success=success, contracts={ fqn: ContractInfoLsp( - abi=info.abi, is_deployable=(len(info.bytecode) > 0) + abi=list(self.info_by_fqn[fqn].abi.values()), + is_deployable=(len(bytecode) > 0), ) - for fqn, info in _compilation.items() + for fqn, bytecode in self.bytecode_by_fqn.items() }, errors={ source_unit_name: [ @@ -502,7 +554,7 @@ def fqn_to_contract_abi(fqn: str): return None, self.info_by_fqn[fqn].abi try: - bytecode = self.compilation[params.contract_fqn].bytecode + bytecode = self.bytecode_by_fqn[params.contract_fqn] for match in LIBRARY_PLACEHOLDER_REGEX.finditer(bytecode): lib_id = bytes.fromhex(match.group(0)[3:-3]) @@ -556,7 +608,9 @@ def fqn_to_contract_abi(fqn: str): data=params.contract_fqn.encode("utf-8"), digest_bits=256 ).digest()[:17] if lib_id in self.libraries: - chain._deployed_libraries[lib_id].append(tx.return_value) + chain._deployed_libraries[lib_id].append( + Library(tx.return_value, chain) + ) return SakeDeployResult( success=success, diff --git a/wake/lsp/server.py b/wake/lsp/server.py index 6709b394..13a7ec2e 100644 --- a/wake/lsp/server.py +++ b/wake/lsp/server.py @@ -137,6 +137,7 @@ SakeGetAbiParams, SakeGetBalancesParams, SakeLoadStateParams, + SakeLoadWorkspaceStateParams, SakeParams, SakeSetBalancesParams, SakeSetLabelParams, @@ -310,6 +311,20 @@ def __init__( ), None, ), + RequestMethodEnum.SAKE_LOAD_WORKSPACE_STATE: ( + lambda params: ( + self.__sake_context.load_workspace_state( # pyright: ignore reportAttributeAccessIssue + params, + ) + ), + SakeLoadWorkspaceStateParams, + ), + RequestMethodEnum.SAKE_SAVE_WORKSPACE_STATE: ( + lambda params: ( + self.__sake_context.save_workspace_state() # pyright: ignore reportAttributeAccessIssue + ), + None, + ), RequestMethodEnum.SAKE_CONNECT_CHAIN: ( lambda params: ( self.__sake_context.connect_chain( # pyright: ignore reportAttributeAccessIssue