State Machine Verification via Rust Traits

This is @jaybosamiya's take on converting the parameterized modules example from Dafny to Rust's traits system (as part of discussion of whether we need parameterized modules in Dust).

// Make sure you are looking at
// `Test/modules/StateMachineRefinement.dfy` from
// https://github.com/secure-foundations/dafny/ on branch
// `mod-params-bryan` (currently at commit ba60b45e4082) too, since I
// regularly refer to it below, and the code here is an (almost)
// direct translation.

// ------- Some boilerplate first -------
//
// Just to make stuff more readable, since I am writing in Rust rather
// than Dust at the moment, I define a bunch of aliases, to make the
// actual code understandable. The proofs themselves are unchecked at
// the moment, since the aim of this file is simply to show what the
// code might look like for writing.

#![allow(dead_code, non_snake_case, non_camel_case_types)]
#[allow(non_camel_case_types)]
type int = i32;
type Map<K, V> = std::collections::HashMap<K, V>;
fn is_updated_map<K, V>(_m1: &Map<K, V>, _m2: &Map<K, V>, _k: &K, _v: &V) -> bool {
    todo!()
}
fn assert(_: bool) {}

// ------- Actual code starts here -------

// We make `TransitionLabel` a direct concrete type in this module. We
// _could_ make it parametric by making everything in this module take
// a type parameter, and then instantiate it at time of use, but that
// seems a bit unweildy.
#[derive(Clone, Copy)]
enum TransitionLabel {
    Query { k: int, value: int },
    Insert { k: int, value: int },
}
use TransitionLabel::*;

// The `Variables` field from the DFY file is "most important" (based
// upon usage here, and the rest of the file), so it is just pulled
// into being `Self` here, since traits (at a high level) are like
// parametric modules that have a single "high priority" type within
// them called `Self`, and allow for nicer syntax in such a
// scenario. Also, since we are using `Self`, the first argument can
// just be made `self`. Since they are read-only, I make them take
// read-only references (this is only specific to Rust; in Dust
// instead, this would be a spec-world definition, as I understand it,
// meaning that we won't need to do this reference thing unless I'm
// mistaken).
trait StateMachine {
    fn predicate_init(&self) -> bool;
    fn predicate_next(&self, next: &Self, l: &TransitionLabel) -> bool;
}

#[derive(Clone, Default, PartialEq, Eq)]
struct MapStateMachine {
    variables: Map<int, int>,
}

impl StateMachine for MapStateMachine {
    fn predicate_init(&self) -> bool {
        self.variables.is_empty()
    }
    fn predicate_next(&self, next: &Self, l: &TransitionLabel) -> bool {
        match l {
            Query { k, value } => {
                self.variables.get(k) == Some(value) && next.variables == self.variables
            }
            Insert { k, value } => is_updated_map(&self.variables, &next.variables, k, value),
        }
    }
}

// Here we define a trait that doesn't have any of its types as the
// "higher importance" `Self`. Instead, this trait is meant to be only
// implemented on a ZST (zero-sized type), which we call
// `Refinement`. This helps us keep `L` and `H` symmetric in usage,
// and also allows for conveniently specifying the refinements
// needed. In some scenarios, we may instead want to make `L` or `H`
// "higher importance" and promote them to being `Self` for this trait
// rather than type parameters. However, we lose the symmetry at that
// point, and based on how the code is used here, it seemed better to
// maintain symmetry.
trait StateMachineRefinement<L: StateMachine, H: StateMachine> {
    fn I(s: L) -> H;

    // In Dust, you'd write out the `requires` and `ensures` clauses
    // here for the lemmas. Since I'm writing in Rust here, I've just
    // added the names in the code, and inserted the requires and
    // ensures clauses as comments.
    fn lemma_init_refinement(s: L)
    // requires s.predicate_init()
    // ensures I(s).predicate_init()
    ;
    fn lemma_next_refinement(s0: L, s1: L, l: TransitionLabel)
    // requires s0.predicate_next(s1, l)
    // ensures I(s).predicate_next(I(s1), l)
    ;
}

// `Refinement` is a zero-sized type (ZST) which means that it is
// isomorphic to `unit`. This means that it stores nothing and takes
// up no space. Instead, it is used only to specify things at a type
// level. Note that (like other types in Rust) just because the
// representation (i.e., zero bytes in this case) is equal, the types
// aren't equal-- this makes it easy to help specify stuff at the type
// level, and get some nice guarantees out of it. It also allows for
// making certain syntax more convenient. For example, here, rather
// than just being an plain ZST, `Refinement` has two arguments (which
// _must_ be state machines). These are maintained to be the same as
// the trait that will be defined on it. This allows us to make it
// more convenient to later refer to the specific lemmas we want,
// rather than needing to perform elaborate trait-selecting type
// syntax (eg: instead of `<Refinement as StateMachineRefinement<X,
// Y>>::lemma_init_refinement`, which would happen if `Refinement`
// took no type parameters, we can simply do the much nicer
// `Refinement<X, Y>::lemma_init_refinement`).
struct Refinement<L: StateMachine, H: StateMachine> {
    _phantom: std::marker::PhantomData<(L, H)>,
}

// Defining `refinement` composition generically (using _only_ type
// parameters) is not quite possible, due to E0203 "A type parameter
// that is specified for impl is not constrained." The core issue with
// doing it only with type parameters is that `M` ends up
// unconstrained, and thus cannot be inferred, and thus Rust disallows
// doing it directly. However, specifying an explicit `M` is
// fine. While we _could_ manually write out code for each such
// explicit `L, M, H`-tuple (and indeed the definition in the Dafny
// `StateMachineRefinement.dfy` does it only on a single
// manually-written tuple, and doesn't generalize things as we do
// here), it is easy enough to just make it into a single easy-to-use
// macro, since the body is going to be the same. The example below
// (after the macro definition) shows the refinement in action.
//
// If you aren't used to `macro_rules!` macros, this is a _very_ basic
// usage of them, where it is only doing a syntactic replacement
// without much else going on. You could easily just (manually) remove
// all the `$` signs, and replace out `L`, `M`, `H` with your own
// types and you'd get the output you are looking for. Rust's macros
// are not textual replacement, and provide stronger guarantees than C
// macros; but in this particular scenario (because it is so simple)
// you can think of it as simple textual replacement.
macro_rules! compose_refinement {
    ($L:ty, $M:ty, $H:ty) => {
        impl StateMachineRefinement<$L, $H> for Refinement<$L, $H>
        where
            $H: StateMachine,
            $L: StateMachine,
            $M: StateMachine,
            Refinement<$L, $M>: StateMachineRefinement<$L, $M>,
            Refinement<$M, $H>: StateMachineRefinement<$M, $H>,
        {
            fn I(s: $L) -> $H {
                <Refinement<$M, $H>>::I(
                    // We need to explicitly tell it which specific refinement
                    // we are picking here for each of the `I`s. Not sure if
                    // we can clean it up / do better.
                    <Refinement<$L, $M>>::I(s),
                )
            }
            fn lemma_init_refinement(s: $L) {
                // We need `clone` here since we are in Rust; in Dust,
                // this would be in proof land, so presumably wouldn't
                // need `clone` since we're working on ghost types at
                // that point
                <Refinement<$L, $M>>::lemma_init_refinement(s.clone());
                <Refinement<$M, $H>>::lemma_init_refinement(<Refinement<$L, $M>>::I(s));
            }
            fn lemma_next_refinement(s0: $L, s1: $L, l: TransitionLabel) {
                <Refinement<$L, $M>>::lemma_next_refinement(s0.clone(), s1.clone(), l);
                <Refinement<$M, $H>>::lemma_next_refinement(
                    <Refinement<$L, $M>>::I(s0),
                    <Refinement<$L, $M>>::I(s1),
                    l,
                );
            }
        }
    };
}

// `MapStateMachine2` and `MapStateMachine3` don't have anything
// special in them; they're just there to show the refinement; feel
// free to jump to next comment directly.
#[derive(Clone, Default, PartialEq, Eq)]
struct MapStateMachine2 {
    x: Map<int, int>,
}

impl StateMachine for MapStateMachine2 {
    fn predicate_init(&self) -> bool {
        self.x.is_empty()
    }
    fn predicate_next(&self, next: &Self, l: &TransitionLabel) -> bool {
        match l {
            Query { k, value } => self.x.get(k) == Some(value) && next.x == self.x,
            Insert { k, value } => is_updated_map(&self.x, &next.x, k, value),
        }
    }
}

#[derive(Clone, Default, PartialEq, Eq)]
struct MapStateMachine3 {
    y: Map<int, int>,
}

impl StateMachine for MapStateMachine3 {
    fn predicate_init(&self) -> bool {
        self.y.is_empty()
    }
    fn predicate_next(&self, next: &Self, l: &TransitionLabel) -> bool {
        match l {
            Query { k, value } => self.y.get(k) == Some(value) && next.y == self.y,
            Insert { k, value } => is_updated_map(&self.y, &next.y, k, value),
        }
    }
}

// Show and prove the base refinements. Because of the specific
// machines picked, the proofs are trivial.
impl StateMachineRefinement<MapStateMachine, MapStateMachine2>
    for Refinement<MapStateMachine, MapStateMachine2>
{
    fn I(s: MapStateMachine) -> MapStateMachine2 {
        MapStateMachine2 { x: s.variables }
    }
    fn lemma_init_refinement(_s: MapStateMachine) {}
    fn lemma_next_refinement(_s0: MapStateMachine, _s1: MapStateMachine, _l: TransitionLabel) {}
}
impl StateMachineRefinement<MapStateMachine2, MapStateMachine3>
    for Refinement<MapStateMachine2, MapStateMachine3>
{
    fn I(s: MapStateMachine2) -> MapStateMachine3 {
        MapStateMachine3 { y: s.x }
    }
    fn lemma_init_refinement(_s: MapStateMachine2) {}
    fn lemma_next_refinement(_s0: MapStateMachine2, _s1: MapStateMachine2, _l: TransitionLabel) {}
}

// Perform the refinement composition using the macro we defined
// earlier. If we had more refinements to be done, it'd just require
// similar lines to provide the transitive proofs.
compose_refinement! {MapStateMachine, MapStateMachine2, MapStateMachine3}

// Why not use a module to group stuff together? This is used for
// testing (from what I can tell in the DFY anyways). Modules in Rust
// only perform grouping and namespacing of code.
mod Final {
    use super::*;
    use std::iter::FromIterator;

    type A = MapStateMachine;
    type B = MapStateMachine3;
    type BigRef = Refinement<A, B>;

    fn lemma_stuff() {
        let s: A = A::default();
        let bigrefis = BigRef::I;
        assert(bigrefis(s.clone()) == B { y: Map::new() });
        BigRef::lemma_init_refinement(s);

        BigRef::lemma_next_refinement(
            A {
                // There is no direct syntax for `map[1 := 2]` in
                // Rust, so we make it from an iterator.
                variables: Map::from_iter(std::iter::once((1, 2))),
            },
            A {
                variables: Map::from_iter(std::iter::once((1, 2))),
            },
            Query { k: 1, value: 2 },
        );
    }

    // lemma `names_for_same_type` is not needed, since we are using a
    // common one throughout the module anyways. If we _were_ making
    // it parametric throughout this file, we'd go and add it as a
    // type parameter to everything, in which case it still turns out
    // to all be equal.
}