Make `Jet` object-safe

[apoelstra]

In several places, most recently BlockstreamResearch/SimplicityHL#224, we've wanted to switch jets at runtime or at least deal with multiple programs whose jet types might differ or be unknown until runtime.

To make this work we need to stop using compile-time generics for the jet. This may mean a Box<dyn Jet> or maybe it means a concrete type. In either case, we need to deal with the transaction environment. Our existing trait has

    /// Obtains a C FFI compatible environment for the jet.
    fn c_jet_env(env: &Self::Environment) -> &Self::CJetEnvironment;

    /// Obtain the FFI C pointer for the jet.
    fn c_jet_ptr(&self) -> &dyn Fn(&mut CFrameItem, CFrameItem, &Self::CJetEnvironment) -> bool;

Reflecting that depending on the type of jets used, the transaction environment should be a different struct (e.g. data for a Bitcoin vs Elements transaction). We need to invert the generics somehow so that Jet can be object-safe, and presumably we would have some sort of Environment object which was generic over the jet.

It may be that we want to use the ghostcell pattern that we use for types::Context (or maybe we even want to combine this with types::Context...though I don't think so because the type context is something that exists only at construction time, while the transaction environment exists for the whole lifecycle of a program. So probably we just want to do runtime checks.

[apoelstra]

A straightforward inversion would be:

trait JetEnvironment (
    type Jet; // maybe we need a Jet trait bound?
    type CJetEnvironment;

    fn c_env(&self) -> &Self::CJetEnvironment;
    // do we actually need a `dyn Fn`? I think a `fn` would suffice.)
    fn c_jet_ptr(&self, jet: &Self::Jet) -> &dyn Fn(&mut CFrameItem, CFrameItem, &Self::CJetEnvironment) -> bool;
)

// The Jet trait now has no knowledge of environments
pub trait Jet:
    Copy + Eq + Ord + Hash + std::fmt::Debug + std::fmt::Display + std::str::FromStr + 'static
{
    /// Return the CMR of the jet.
    fn cmr(&self) -> Cmr;

    /// Return the source type of the jet.
    fn source_ty(&self) -> TypeName;

    /// Return the target type of the jet.
    fn target_ty(&self) -> TypeName;

    /// Encode the jet to bits.
    fn encode<W: Write>(&self, w: &mut BitWriter<W>) -> std::io::Result<usize>;

    /// Decode a jet from bits.
    fn decode<I: Iterator<Item = u8>>(bits: &mut BitIter<I>) -> Result<Self, decode::Error>;

    /// Return the cost of the jet.
    fn cost(&self) -> Cost;

This actually doesn't get us all the way to object-safety because we still have the decode method which returns Self. It may be that we just want to move the decode method off of Jet and into JetEnvironment. (And then we could put a JetEnvironment into types::Context, and make this be responsible for constructing jets, which would go some way toward preventing users from mixing different jet types during construction.)

[ivanlele]

I was playing with this today to see in what form it might work. Rust’s rules about what can be placed inside Box<dyn T> are quite strict (https://doc.rust-lang.org/reference/items/traits.html#dyn-compatibility).

Some Jet supertraits are not object-safe, so I was forced to came up with a workaround. Also encode needs to be constrained to a specific type, so I thought to use &mut BitWriter<&mut Vec> (Don't know if this will be good ideas, as this also raices a lot of work to do).

With this, I started updating the affected areas and later attempted to change the Node types. After a few hours, I noticed that the changes had already grown to several thousand lines. Even my generic hell stopped to look that bad, and there were still more changes to come, since everything depends on J: Jet, and we also changed decoding logic.

I was thinking about whether this could be implemented in several small PRs. Perhaps we could gradually move functionality from the Jet trait into JetEnvironment. Doing everything in a single step forces us to update a large amount of related code. Another think we maybe can create a separate version and switch to it at some point.

I was also considering when types::Context should come into play. Would it be better to pass it directly into the decode function, or should it instead be exposed via a separate ctx() method on JetEnvironment?

pub trait JetEnvironment {
    type Jet: Jet;
    type CJetEnvironment;

    /// Obtains a C FFI compatible environment for the jet.
    fn c_jet_env(&self) -> &Self::CJetEnvironment;

    /// Decode a jet from bits.
    fn decode_jet<I: Iterator<Item = u8>>(
        ctx: types::Context<'_>,
        bits: &mut BitIter<I>,
    ) -> Result<Box<dyn Jet>, decode::Error>;

    /// Obtain the FFI C pointer for the jet.
    fn c_jet_ptr(
        jet: &Self::Jet,
    ) -> fn(&mut CFrameItem, CFrameItem, &Self::CJetEnvironment) -> bool;
}

/// Helper trait to enable `Clone`, `Hash`, `Eq`, and `Ord` for `Box<dyn Jet>`.
trait DynJet {
    fn dyn_clone(&self) -> Box<dyn Jet>;
    fn dyn_hash(&self, state: &mut dyn Hasher);
    fn dyn_eq(&self, other: &dyn Jet) -> bool;
    fn dyn_cmp(&self, other: &dyn Jet) -> std::cmp::Ordering;
}

impl<T: Clone + Hash + Jet> DynJet for T {
    fn dyn_clone(&self) -> Box<dyn Jet> {
        Box::new(self.clone())
    }

    fn dyn_hash(&self, mut state: &mut dyn Hasher) {
        self.hash(&mut state)
    }

    fn dyn_eq(&self, other: &dyn Jet) -> bool {
        self.cmr() == other.cmr()
    }

    fn dyn_cmp(&self, other: &dyn Jet) -> std::cmp::Ordering {
        self.cmr().cmp(&other.cmr())
    }
}

impl Clone for Box<dyn Jet> {
    fn clone(&self) -> Self {
        (**self).dyn_clone()
    }
}

impl Hash for Box<dyn Jet> {
    fn hash<H: Hasher>(&self, state: &mut H) {
        (**self).dyn_hash(state)
    }
}

impl PartialEq for Box<dyn Jet> {
    fn eq(&self, other: &Self) -> bool {
        (**self).dyn_eq(&**other)
    }
}

impl Eq for Box<dyn Jet> {}

impl PartialOrd for Box<dyn Jet> {
    fn partial_cmp(&self, other: &Self) -> Option<std::cmp::Ordering> {
        Some(self.cmp(other))
    }
}

impl Ord for Box<dyn Jet> {
    fn cmp(&self, other: &Self) -> std::cmp::Ordering {
        (**self).dyn_cmp(&**other)
    }
}

/// Family of jets that share an encoding scheme and execution environment.
///
/// Jets are single nodes that read an input,
/// process it internally using foreign C code _(black box)_,
/// and produce an output.
/// Jets may read values from their _environment_.
///
/// Jets are **always** leaves in a Simplicity DAG.
pub trait Jet: DynJet + std::fmt::Display + std::fmt::Debug + 'static {
    /// Return the CMR of the jet.
    fn cmr(&self) -> Cmr;

    /// Return the source type of the jet.
    fn source_ty(&self) -> TypeName;

    /// Return the target type of the jet.
    fn target_ty(&self) -> TypeName;

    /// Encode the jet to bits.
    fn encode(&self, w: &mut BitWriter<&mut Vec<u8>>) -> std::io::Result<usize>;

    /// Return the cost of the jet.
    fn cost(&self) -> Cost;
}

[apoelstra]

I was thinking about whether this could be implemented in several small PRs. Perhaps we could gradually move functionality from the Jet trait into JetEnvironment.

Yeah. So, we know that we've gotta move the c_jet_env and c_jet_ptr methods, so we could move those one by one in separate PRs, which shouldn't be too noisy since they're not super common methods. Creating and using the JetEnvironment trait might also be its own PR, even if it's initially an empty trait. (Unsure.)

As you observe, some supertraits are not object-safe (I see FromStr but probably there are others) so we would also need to e.g. replace FromStr with a dedicated parse method or something. This can also be its own PR.

I was also considering when types::Context should come into play.

Let's try to do the "easy" stuff like the above, and then once we've gotten the trait to be object-safe, we can re-assess how to use it.

With this, I started updating the affected areas and later attempted to change the Node types. After a few hours, I noticed that the changes had already grown to several thousand lines

So, unfortunately changing Node is going to be a hell of a change. I hope we can remove the generic and replace Jet(J) with Jet(Box<dyn Jet>) and make no other changes in a single PR, since that PR will have a huge diff but hopefully be mechanical and simple. Then once the dust clears from having removed the generic, we can try to use the jet type "properly" by pulling types::Context in and doing compatibility checks and so on.

[ivanlele]

Yeah. So, we know that we've gotta move the c_jet_env and c_jet_ptr methods, so we could move those one by one in separate PRs, which shouldn't be too noisy since they're not super common methods. Creating and using the JetEnvironment trait might also be its own PR, even if it's initially an empty trait. (Unsure.)

I made #351 which should cover all these.