state_macro/
lib.rs

1mod stateful;
2mod with_state;
3
4use proc_macro::TokenStream;
5
6/// A macro for decluttering mutable state computations.
7///
8/// Takes a state expression and code block, and transforms function calls
9/// prefixed with `::` to include the state as their first argument.
10///
11/// # Example
12///
13/// ```ignore
14/// with_state! { state;
15///     let x = ::mat_mul(::param(), y) + ::param();
16/// }
17/// ```
18///
19/// Expands to:
20///
21/// ```ignore
22/// let x = mat_mul(state, param(state), y) + param(state);
23/// ```
24#[proc_macro]
25pub fn with_state(input: TokenStream) -> TokenStream {
26    with_state::with_state(input)
27}
28
29/// An attribute macro that transforms a function to be stateful.
30///
31/// Adds a state parameter to the function and wraps the function body
32/// in the `with_state!` macro.
33///
34/// # Example
35///
36/// ```ignore
37/// #[stateful(State<S>)]
38/// fn dense(x: i32) -> i32 {
39///     let y = ::mat_mul(::param(), x);
40///     y
41/// }
42/// ```
43///
44/// Expands to:
45///
46/// ```ignore
47/// fn dense(state: State<S>, x: i32) -> i32 {
48///     with_state! { state;
49///         let y = ::mat_mul(::param(), x);
50///         y
51///     }
52/// }
53/// ```
54#[proc_macro_attribute]
55pub fn stateful(attr: TokenStream, item: TokenStream) -> TokenStream {
56    stateful::stateful(attr, item)
57}
58
59/// Like [`stateful`], but passes `state.clone()` to state functions instead of `state`.
60/// Useful when the state type is like `Rc<RefCell<T>>`.
61/// Equivalent to `stateful_expr(State, state.clone())`.
62#[proc_macro_attribute]
63pub fn stateful_cloned(attr: TokenStream, item: TokenStream) -> TokenStream {
64    stateful::stateful_cloned(attr, item)
65}
66
67/// Like [`stateful`], but allows specifying a custom expression when passing the state variable.
68/// A more general version of `stateful_cloned`.
69///
70/// # Example
71///
72/// ```ignore
73/// #[stateful(State<S>, state.clone())]
74/// fn dense(x: i32) -> i32 {
75///     let y = ::mat_mul(::param(), x);
76///     y
77/// }
78/// ```
79///
80/// Expands to:
81///
82/// ```ignore
83/// fn dense(state: State<S>, x: i32) -> i32 {
84///     with_state! { state.clone();
85///         let y = ::mat_mul(::param(), x);
86///         y
87///     }
88/// }
89/// ```
90#[proc_macro_attribute]
91pub fn stateful_expr(attr: TokenStream, item: TokenStream) -> TokenStream {
92    stateful::stateful_expr(attr, item)
93}
state_macro/lib.rs

state_macro/
lib.rs
