Skip to content

Commit

Permalink
add python bindings
Browse files Browse the repository at this point in the history
  • Loading branch information
arvidn committed Sep 24, 2023
1 parent ca1dedd commit 5afc328
Show file tree
Hide file tree
Showing 5 changed files with 100 additions and 0 deletions.
10 changes: 10 additions & 0 deletions src/error.rs
Original file line number Diff line number Diff line change
Expand Up @@ -2,6 +2,9 @@ use crate::gen::validation_error::ValidationErr;
use clvmr::reduction::EvalErr;
use thiserror::Error;

#[cfg(feature = "py-bindings")]
use pyo3::PyErr;

#[derive(Debug, Clone, PartialEq, Eq, Error)]
pub enum Error {
#[error("CLVM {0}")]
Expand Down Expand Up @@ -40,4 +43,11 @@ pub enum Error {
Custom(String),
}

#[cfg(feature = "py-bindings")]
impl std::convert::From<Error> for PyErr {
fn from(err: Error) -> PyErr {
pyo3::exceptions::PyValueError::new_err(err.to_string())
}
}

pub type Result<T> = std::result::Result<T, Error>;
23 changes: 23 additions & 0 deletions tests/test_fast_forward.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,23 @@
from chia_rs import Spend, CoinSpend, Coin, supports_fast_forward, fast_forward_singleton
import pytest

@pytest.mark.parametrize("file", ["bb13", "e3c0"])
def test_supports_fast_forward(file: str) -> None:
with open(f"ff-tests/{file}.spend", "rb") as f:
spend = CoinSpend.from_bytes(f.read())
assert supports_fast_forward(spend)

@pytest.mark.parametrize("file", ["bb13", "e3c0"])
def test_fast_forward_singleton(file: str) -> None:
with open(f"ff-tests/{file}.spend", "rb") as f:
spend = CoinSpend.from_bytes(f.read())

parents_parent = bytes([0] * 32)
new_parent = Coin(parents_parent, spend.coin.puzzle_hash, spend.coin.amount)
new_coin = Coin(new_parent.name(), new_parent.puzzle_hash, new_parent.amount)
new_solution = fast_forward_singleton(spend, new_coin, new_parent)

expected = bytearray(bytes(spend.solution))
# this is where the parent's parent coin ID lives in the solution
expected[3:35] = parents_parent
assert expected == new_solution
3 changes: 3 additions & 0 deletions wheel/chia_rs.pyi
Original file line number Diff line number Diff line change
Expand Up @@ -13,6 +13,9 @@ def solution_generator_backrefs(spends: Sequence[Tuple[Coin, bytes, bytes]]) ->

def compute_merkle_set_root(items: Sequence[bytes]) -> bytes: ...

def supports_fast_forward(spend: CoinSpend) -> bool : ...
def fast_forward_singleton(spend: CoinSpend, new_coin: Coin, new_parent: Coin) -> bytes: ...

def run_block_generator(
program: ReadableBuffer, args: List[ReadableBuffer], max_cost: int, flags: int
) -> Tuple[Optional[int], Optional[SpendBundleConditions]]: ...
Expand Down
3 changes: 3 additions & 0 deletions wheel/generate_type_stubs.py
Original file line number Diff line number Diff line change
Expand Up @@ -164,6 +164,9 @@ def solution_generator_backrefs(spends: Sequence[Tuple[Coin, bytes, bytes]]) ->
def compute_merkle_set_root(items: Sequence[bytes]) -> bytes: ...
def supports_fast_forward(spend: CoinSpend) -> bool : ...
def fast_forward_singleton(spend: CoinSpend, new_coin: Coin, new_parent: Coin) -> bytes: ...
def run_block_generator(
program: ReadableBuffer, args: List[ReadableBuffer], max_cost: int, flags: int
) -> Tuple[Optional[int], Optional[SpendBundleConditions]]: ...
Expand Down
61 changes: 61 additions & 0 deletions wheel/src/api.rs
Original file line number Diff line number Diff line change
Expand Up @@ -50,6 +50,7 @@ use std::iter::zip;
use crate::run_program::{run_chia_program, serialized_length};

use crate::adapt_response::eval_err_to_pyresult;
use chia::fast_forward::fast_forward_singleton as native_ff;
use chia::gen::get_puzzle_and_solution::get_puzzle_and_solution_for_coin as parse_puzzle_solution;
use chia::gen::validation_error::ValidationErr;
use clvmr::allocator::NodePtr;
Expand Down Expand Up @@ -265,6 +266,64 @@ impl AugSchemeMPL {
}
}

#[pyfunction]
fn supports_fast_forward(spend: &CoinSpend) -> bool {
// the test function just attempts the rebase onto a dummy parent coin
let new_parent = Coin {
parent_coin_info: [0_u8; 32].into(),
puzzle_hash: spend.coin.puzzle_hash,
amount: spend.coin.amount,
};
let new_coin = Coin {
parent_coin_info: new_parent.coin_id().into(),
puzzle_hash: spend.coin.puzzle_hash,
amount: spend.coin.amount,
};

let mut a = make_allocator(LIMIT_HEAP);
let Ok(puzzle) = node_from_bytes(&mut a, spend.puzzle_reveal.as_slice()) else {
return false;
};
let Ok(solution) = node_from_bytes(&mut a, spend.solution.as_slice()) else {
return false;
};

native_ff(
&mut a,
puzzle,
solution,
&spend.coin,
&new_coin,
&new_parent,
)
.is_ok()
}

#[pyfunction]
fn fast_forward_singleton<'p>(
py: Python<'p>,
spend: &CoinSpend,
new_coin: &Coin,
new_parent: &Coin,
) -> PyResult<&'p PyBytes> {
let mut a = make_allocator(LIMIT_HEAP);
let puzzle = node_from_bytes(&mut a, spend.puzzle_reveal.as_slice())?;
let solution = node_from_bytes(&mut a, spend.solution.as_slice())?;

let new_solution = native_ff(
&mut a,
puzzle,
solution,
&spend.coin,
new_coin,
new_parent,
)?;
Ok(PyBytes::new(
py,
node_to_bytes(&a, new_solution)?.as_slice(),
))
}

#[pymodule]
pub fn chia_rs(py: Python, m: &PyModule) -> PyResult<()> {
// generator functions
Expand All @@ -273,6 +332,8 @@ pub fn chia_rs(py: Python, m: &PyModule) -> PyResult<()> {
m.add_function(wrap_pyfunction!(run_puzzle, m)?)?;
m.add_function(wrap_pyfunction!(solution_generator, m)?)?;
m.add_function(wrap_pyfunction!(solution_generator_backrefs, m)?)?;
m.add_function(wrap_pyfunction!(supports_fast_forward, m)?)?;
m.add_function(wrap_pyfunction!(fast_forward_singleton, m)?)?;
m.add_class::<PySpendBundleConditions>()?;
m.add(
"ELIGIBLE_FOR_DEDUP",
Expand Down

0 comments on commit 5afc328

Please sign in to comment.