Skip to content

Instantly share code, notes, and snippets.

@byte-sourcerer
Last active March 12, 2023 13:40
Show Gist options
  • Save byte-sourcerer/07f4fba65b2b88c5a9c55aa180f0f34d to your computer and use it in GitHub Desktop.
Save byte-sourcerer/07f4fba65b2b88c5a9c55aa180f0f34d to your computer and use it in GitHub Desktop.
Permission Pattern for Rust
/// ref: [zhihu](https://zhuanlan.zhihu.com/p/378252833)
use crate::foo::{IncreaseOnly, new_increase_only};
fn main() {
let mut x1 = new_increase_only!(0);
let mut x2 = new_increase_only!(2);
let y = 1;
// only can be used in `x1`
let le = x1.check_less_or_equal(y).unwrap();
// ERROR! expected struct `main::__Permit`, found a different struct `main::__Permit`
x2.update(le);
}
mod foo {
use std::marker::PhantomData;
pub struct LessOrEqual<Permit> {
x: u64,
permission: PhantomData<*mut Permit>,
}
impl<Permit> LessOrEqual<Permit> {
fn new(x: u64) -> Self {
Self { x, permission: PhantomData }
}
}
#[repr(transparent)]
pub struct IncreaseOnly<Permit> {
x: u64,
permission: PhantomData<*mut Permit>,
}
macro_rules! new_increase_only {
($value: expr) => {
{
struct __Permit;
#[repr(transparent)]
struct RawIncreaseOnly<Permit> {
x: u64,
permission: std::marker::PhantomData<*mut Permit>,
}
// we cannot use `IncreaseOnly::new` here
// since it will expose private API and lead to compile error
unsafe {
std::mem::transmute::<RawIncreaseOnly<__Permit>, IncreaseOnly<__Permit>>(
RawIncreaseOnly { x: $value, permission: std::marker::PhantomData }
)
}
}
}
}
pub(crate) use new_increase_only;
impl<Permit> IncreaseOnly<Permit> {
fn new(x: u64) -> Self {
Self { x, permission: PhantomData }
}
pub fn check_less_or_equal(&self, other: u64) -> Option<LessOrEqual<Permit>> {
if self.x <= other {
Some(LessOrEqual::new(other))
} else {
None
}
}
pub fn update(&mut self, greater_than_self: LessOrEqual<Permit>) {
self.x = greater_than_self.x;
}
}
}
@byte-sourcerer
Copy link
Author

byte-sourcerer commented Mar 3, 2023

Generalization:

use crate::lock::{new_lock, Lock, Unlock};

fn main() {
    let mut l1: Lock<IncreaseOnly, _> = new_lock!(IncreaseOnly(0));
    let mut l2: Lock<IncreaseOnly, _> = new_lock!(IncreaseOnly(0));
    let key1 = l1.get_key(1).unwrap();

    // ERROR
    // l2.unlock(key1);

    l1.unlock(key1);
}

struct IncreaseOnly(u64);

impl Unlock for IncreaseOnly {
    type Args = u64;
    type KeyInner = u64;
    type Effect<T> = Option<T>;
    type Unlock = ();

    fn generate_key(&self, args: Self::Args) -> Self::Effect<Self::KeyInner> {
        if self.0 <= args {
            Some(args)
        } else {
            None
        }
    }

    fn unlock(&mut self, key: Self::KeyInner) -> Self::Unlock {
        self.0 = key;
    }

    fn map_effect<T, U, F>(effect: Self::Effect<T>, f: F) -> Self::Effect<U>
    where
        F: FnOnce(T) -> U,
    {
        effect.map(f)
    }
}

mod lock {
    use std::marker::PhantomData;

    pub type InvariantType<T> = PhantomData<*mut T>;

    #[repr(transparent)]
    pub struct Lock<T, Permit> {
        value: T,
        permission: InvariantType<Permit>,
    }

    impl<T, Permit> Lock<T, Permit>
    where
        T: Unlock,
    {
        pub fn get_key(&self, args: T::Args) -> T::Effect<Key<T::KeyInner, Permit>> {
            T::map_effect(self.value.generate_key(args), Key::new)
        }

        pub fn unlock(&mut self, key: Key<T::KeyInner, Permit>) -> T::Unlock {
            self.value.unlock(key.value)
        }
    }

    pub struct Key<T, Permit> {
        value: T,
        permission: InvariantType<Permit>,
    }

    impl<T, Permit> Key<T, Permit> {
        fn new(value: T) -> Self {
            Self {
                value,
                permission: PhantomData,
            }
        }
    }

    pub trait Unlock {
        type Args;
        type KeyInner;
        type Effect<T>;

        type Unlock;

        fn generate_key(&self, args: Self::Args) -> Self::Effect<Self::KeyInner>;

        fn unlock(&mut self, key: Self::KeyInner) -> Self::Unlock;

        /// `Self::Effect` should be a functor
        fn map_effect<T, U, F>(effect: Self::Effect<T>, f: F) -> Self::Effect<U>
        where
            F: FnOnce(T) -> U;
    }

    macro_rules! new_lock {
        ($lockable: expr) => {{
            struct __Permit;

            #[repr(transparent)]
            pub struct __RawLock<T, Permit> {
                #[allow(dead_code)]
                value: T,
                permission: lock::InvariantType<Permit>,
            }

            unsafe {
                std::mem::transmute::<__RawLock<_, __Permit>, lock::Lock<_, __Permit>>(__RawLock {
                    value: $lockable,
                    permission: std::marker::PhantomData,
                })
            }
        }};
    }

    pub(crate) use new_lock;
}

@byte-sourcerer
Copy link
Author

Update:

use crate::db::Db;
use crate::id::make_id;

mod db {
    use std::collections::HashMap;
    use crate::id::{Id, make_id};

    pub struct Db<Brand> {
        id: Id<Brand>,
        data: HashMap<String, i32>
    }

    impl<Brand> Db<Brand> {
        pub fn new(id: Id<Brand>) -> Self {
            Self {
                id,
                data: HashMap::new(),
            }
        }

        pub fn insert(&mut self, key: String, value: i32) -> Permit<Brand> {
            self.data.insert(key, value);
            Permit::new(self.id.clone())
        }

        pub fn get(&self, key: &str, _permit: Permit<Brand>) -> Option<&i32> {
            self.data.get(key)
        }
    }

    pub struct Permit<Brand>(Id<Brand>);

    impl<Brand> Permit<Brand> {
        fn new(id: Id<Brand>) -> Self {
            Self(id)
        }
    }
}


mod id {
    use std::marker::PhantomData;

    pub type InvariantType<T> = PhantomData<*mut T>;

    #[repr(transparent)]
    #[derive(Eq, PartialEq)]
    pub struct Id<Brand>(InvariantType<Brand>);

    impl<Brand> Id<Brand> {
        fn new() -> Self {
            Self(PhantomData)
        }
    }

    impl<Brand> Clone for Id<Brand> {
        fn clone(&self) -> Self {
            Id::new()
        }
    }

    macro_rules! make_id {
        () => {
            {
                use crate::id::InvariantType;
                use std::marker::PhantomData;
                use crate::id::Id;

                struct __Brand;

                #[repr(transparent)]
                pub struct RawId<Brand>(InvariantType<Brand>);
                let raw_id: RawId<__Brand> = RawId(PhantomData);

                unsafe {
                    std::mem::transmute::<RawId<__Brand>, Id<__Brand>>(raw_id)
                }
            }
        };
    }

    pub(crate) use make_id;
}


fn main() {
    let id1 = make_id!();
    let id2 = make_id!();
    let mut db1 = Db::new(id1);
    let db2 = Db::new(id2);
    let key = String::from("key");
    let permit1 = db1.insert(key.clone(), 1);
    assert_eq!(db1.get(&key, permit1), Some(&1));
    db2.get(&key, permit1); // ERROR
}

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment