From 9858b27cc103e483687bcfc769057345c6c5edb7 Mon Sep 17 00:00:00 2001 From: mhh Date: Sun, 17 Mar 2024 14:30:12 +0100 Subject: [PATCH] Add new instructions for the Compute Budget program to interact with the compute budget program. --- src/spl/compute_budget/__init__.py | 1 + src/spl/compute_budget/_layouts.py | 42 ++++ src/spl/compute_budget/constants.py | 6 + src/spl/compute_budget/instructions.py | 264 +++++++++++++++++++++++++ 4 files changed, 313 insertions(+) create mode 100644 src/spl/compute_budget/__init__.py create mode 100644 src/spl/compute_budget/_layouts.py create mode 100644 src/spl/compute_budget/constants.py create mode 100644 src/spl/compute_budget/instructions.py diff --git a/src/spl/compute_budget/__init__.py b/src/spl/compute_budget/__init__.py new file mode 100644 index 00000000..6f368f66 --- /dev/null +++ b/src/spl/compute_budget/__init__.py @@ -0,0 +1 @@ +"""Client code for interacting with the compute budget program.""" diff --git a/src/spl/compute_budget/_layouts.py b/src/spl/compute_budget/_layouts.py new file mode 100644 index 00000000..3395339d --- /dev/null +++ b/src/spl/compute_budget/_layouts.py @@ -0,0 +1,42 @@ +"""Token instruction layouts.""" +from enum import IntEnum + +from construct import Bytes, Int8ul, Int32ul, Int64ul, Switch +from construct import Struct as cStruct + +PUBLIC_KEY_LAYOUT = Bytes(32) + + +class InstructionType(IntEnum): + """Token instruction types.""" + + REQUEST_UNITS = 0 + REQUEST_HEAP_FRAME = 1 + SET_COMPUTE_UNIT_LIMIT = 2 + SET_COMPUTE_UNIT_PRICE = 3 + + +_REQUEST_UNITS_LAYOUT = cStruct( + "units" / Int32ul, + "additional_fee" / Int32ul, +) + +_REQUEST_HEAP_FRAME_LAYOUT = cStruct("bytes" / Int32ul) + +_SET_COMPUTE_UNIT_LIMIT_LAYOUT = cStruct("units" / Int32ul) + +_SET_COMPUTE_UNIT_PRICE_LAYOUT = cStruct("micro_lamports" / Int64ul) + +INSTRUCTIONS_LAYOUT = cStruct( + "instruction_type" / Int8ul, + "args" + / Switch( + lambda this: this.instruction_type, + { + InstructionType.REQUEST_UNITS: _REQUEST_UNITS_LAYOUT, + InstructionType.REQUEST_HEAP_FRAME: _REQUEST_HEAP_FRAME_LAYOUT, + InstructionType.SET_COMPUTE_UNIT_LIMIT: _SET_COMPUTE_UNIT_LIMIT_LAYOUT, + InstructionType.SET_COMPUTE_UNIT_PRICE: _SET_COMPUTE_UNIT_PRICE_LAYOUT, + }, + ), +) diff --git a/src/spl/compute_budget/constants.py b/src/spl/compute_budget/constants.py new file mode 100644 index 00000000..d979084d --- /dev/null +++ b/src/spl/compute_budget/constants.py @@ -0,0 +1,6 @@ +"""Compute Budget token constants.""" + +from solders.pubkey import Pubkey + +COMPUTE_BUDGET_PROGRAM_ID: Pubkey = Pubkey.from_string("ComputeBudget111111111111111111111111111111") +"""Public key that identifies the Compute Budget program.""" diff --git a/src/spl/compute_budget/instructions.py b/src/spl/compute_budget/instructions.py new file mode 100644 index 00000000..393de76a --- /dev/null +++ b/src/spl/compute_budget/instructions.py @@ -0,0 +1,264 @@ +"""Instructions for the Compute Budget program.""" +from typing import NamedTuple, Any + +from solders.instruction import Instruction +from typing_extensions import deprecated + +from solana.utils.validate import validate_instruction_keys, validate_instruction_type +from spl.compute_budget.constants import COMPUTE_BUDGET_PROGRAM_ID +from spl.compute_budget._layouts import ( + InstructionType, + INSTRUCTIONS_LAYOUT, +) + +""" + +class InstructionType(IntEnum): + REQUEST_UNITS = 0 + REQUEST_HEAP_FRAME = 1 + SET_COMPUTE_UNIT_LIMIT = 2 + SET_COMPUTE_UNIT_PRICE = 3 + +_REQUEST_UNITS_LAYOUT = cStruct( + "instruction" / Int8ul, + "units" / Int32ul, + "additional_fee" / Int32ul, +) + +_REQUEST_HEAP_FRAME_LAYOUT = cStruct("instruction" / Int8ul, "bytes" / Int32ul) + +_SET_COMPUTE_UNIT_LIMIT_LAYOUT = cStruct("instruction" / Int8ul, "units" / Int32ul) + +_SET_COMPUTE_UNIT_PRICE_LAYOUT = cStruct("instruction" / Int8ul, "micro_lamports" / Int64ul) + +INSTRUCTIONS_LAYOUT = cStruct( + "instruction_type" / Int8ul, + "args" + / Switch( + lambda this: this.instruction_type, + { + InstructionType.REQUEST_UNITS: _REQUEST_UNITS_LAYOUT, + InstructionType.REQUEST_HEAP_FRAME: _REQUEST_HEAP_FRAME_LAYOUT, + InstructionType.SET_COMPUTE_UNIT_LIMIT: _SET_COMPUTE_UNIT_LIMIT_LAYOUT, + InstructionType.SET_COMPUTE_UNIT_PRICE: _SET_COMPUTE_UNIT_PRICE_LAYOUT, + }, + ), +) + +""" + + +@deprecated("Use SetComputeUnitLimitParams instead.") +class RequestUnitsParams(NamedTuple): + """Request units transaction params.""" + + units: int + """Units.""" + additional_fee: int + """Additional fee.""" + + +class RequestHeapFrameParams(NamedTuple): + """Request heap frame transaction params.""" + + bytes: int # noqa: A003 + """The amount of heap frame bytes to additionally request. 32K = 8 Compute Units.""" + + +class SetComputeUnitLimitParams(NamedTuple): + """Set compute unit limit transaction params.""" + + units: int + """Maximum compute units.""" + + +class SetComputeUnitPriceParams(NamedTuple): + """Set compute unit price transaction params.""" + + micro_lamports: int + """Priority fee of one compute unit in micro Lamports. 1 micro Lamport = 0.000001 Lamport = 10^-15 SOL.""" + + +def __parse_and_validate_instruction( + instruction: Instruction, + expected_keys: int, + expected_type: InstructionType, +) -> Any: # Returns a Construct container. + validate_instruction_keys(instruction, expected_keys) + data = INSTRUCTIONS_LAYOUT.parse(instruction.data) + validate_instruction_type(data, expected_type) + return data + + +@deprecated("Use SetComputeUnitLimit instruction instead.") +def decode_request_units(instruction: Instruction) -> RequestUnitsParams: + """Decode a request_units instruction and retrieve the instruction params. + + Args: + instruction: The instruction to decode. + + Returns: + The decoded instruction. + """ + parsed_data = __parse_and_validate_instruction(instruction, 0, InstructionType.REQUEST_UNITS) + return RequestUnitsParams( + units=parsed_data.args.units, + additional_fee=parsed_data.args.additional_fee, + ) + + +def decode_request_heap_frame(instruction: Instruction) -> RequestHeapFrameParams: + """Decode a request_heap_frame instruction and retrieve the instruction params. + + Args: + instruction: The instruction to decode. + + Returns: + The decoded instruction. + """ + parsed_data = __parse_and_validate_instruction(instruction, 0, InstructionType.REQUEST_HEAP_FRAME) + return RequestHeapFrameParams( + bytes=parsed_data.args.bytes, + ) + + +def decode_set_compute_unit_limit(instruction: Instruction) -> SetComputeUnitLimitParams: + """Decode a set_compute_unit_limit instruction and retrieve the instruction params. + + Args: + instruction: The instruction to decode. + + Returns: + The decoded instruction. + """ + parsed_data = __parse_and_validate_instruction(instruction, 0, InstructionType.SET_COMPUTE_UNIT_LIMIT) + return SetComputeUnitLimitParams( + units=parsed_data.args.units, + ) + + +def decode_set_compute_unit_price(instruction: Instruction) -> SetComputeUnitPriceParams: + """Decode a set_compute_unit_price instruction and retrieve the instruction params. + + Args: + instruction: The instruction to decode. + + Returns: + The decoded instruction. + """ + parsed_data = __parse_and_validate_instruction(instruction, 0, InstructionType.SET_COMPUTE_UNIT_PRICE) + return SetComputeUnitPriceParams( + micro_lamports=parsed_data.args.micro_lamports, + ) + + +@deprecated("Use set_compute_unit_limit instead.") +def request_units(params: RequestUnitsParams) -> Instruction: + """Creates a transaction instruction that requests units. + + Example: + >>> params = RequestUnitsParams(units=150_000, additional_fee=0) + >>> type(request_units(params)) + + + Returns: + The instruction to request units. + """ + data = INSTRUCTIONS_LAYOUT.build( + { + "instruction_type": InstructionType.REQUEST_UNITS, + "args": { + "units": params.units, + "additional_fee": params.additional_fee, + }, + } + ) + return Instruction( + accounts=[], + program_id=COMPUTE_BUDGET_PROGRAM_ID, + data=data, + ) + + +def request_heap_frame(params: RequestHeapFrameParams) -> Instruction: + """Creates a transaction instruction that requests heap frame. + + Example: + >>> params = RequestHeapFrameParams(bytes=32_000 * 100) # 100 * 32K = 800 Compute Units + >>> type(request_heap_frame(params)) + + + Returns: + The instruction to request heap frame. + """ + data = INSTRUCTIONS_LAYOUT.build( + { + "instruction_type": InstructionType.REQUEST_HEAP_FRAME, + "args": { + "bytes": params.bytes, + }, + } + ) + return Instruction( + accounts=[], + program_id=COMPUTE_BUDGET_PROGRAM_ID, + data=data, + ) + + +def set_compute_unit_limit(params: SetComputeUnitLimitParams) -> Instruction: + """Creates a transaction instruction that sets the compute unit limit. + + By default, the compute budget is the product of 200,000 Compute Units (CU) * number of instructions, + with a max of 1.4M CU. + + Example: + >>> params = SetComputeUnitLimitParams(units=1_000_000) + >>> type(set_compute_unit_limit(params)) + + + Returns: + The instruction to set the compute unit limit. + """ + data = INSTRUCTIONS_LAYOUT.build( + { + "instruction_type": InstructionType.SET_COMPUTE_UNIT_LIMIT, + "args": { + "units": params.units, + }, + } + ) + return Instruction( + accounts=[], + program_id=COMPUTE_BUDGET_PROGRAM_ID, + data=data, + ) + + +def set_compute_unit_price(params: SetComputeUnitPriceParams) -> Instruction: + """Creates a transaction instruction that sets the compute unit price. + + For example: 1000 micro Lamports = 10^-12 SOL, so for an instruction that uses 150_000 compute units, + the additional priority fee would be 150_000 * 1000 = 150_000_000 micro Lamports = 150 Lamports = 150 * 10^-9 SOL. + + Example: + >>> params = SetComputeUnitPriceParams(micro_lamports=1_000) + >>> type(set_compute_unit_price(params)) + + + Returns: + The instruction to set the compute unit price. + """ + data = INSTRUCTIONS_LAYOUT.build( + { + "instruction_type": InstructionType.SET_COMPUTE_UNIT_PRICE, + "args": { + "micro_lamports": params.micro_lamports, + }, + } + ) + return Instruction( + accounts=[], + program_id=COMPUTE_BUDGET_PROGRAM_ID, + data=data, + )