Skip to main content

shopify_function_provider/
lib.rs

1mod alloc;
2pub mod log;
3pub mod read;
4mod string_interner;
5pub mod write;
6
7use bumpalo::Bump;
8use rmp::encode::ByteBuf;
9use std::cell::RefCell;
10use string_interner::StringInterner;
11use write::State;
12
13pub const PROVIDER_MODULE_NAME: &str =
14    concat!("shopify_function_v", env!("CARGO_PKG_VERSION_MAJOR"));
15
16#[cfg(target_pointer_width = "64")]
17type DoubleUsize = u128;
18#[cfg(target_pointer_width = "32")]
19type DoubleUsize = u64;
20
21struct Context {
22    bump_allocator: bumpalo::Bump,
23    input_bytes: Vec<u8>,
24    output_bytes: ByteBuf,
25    logs: Logs,
26    write_state: State,
27    write_parent_state_stack: Vec<State>,
28    string_interner: StringInterner,
29}
30
31thread_local! {
32    static CONTEXT: RefCell<Context> = RefCell::new(Context::default())
33}
34
35#[cfg(target_family = "wasm")]
36thread_local! {
37    static OUTPUT_AND_LOG_PTRS: RefCell<[usize; 6]> = const { RefCell::new([0; 6]) };
38}
39
40impl Default for Context {
41    fn default() -> Self {
42        Self {
43            bump_allocator: Bump::new(),
44            input_bytes: Vec::new(),
45            output_bytes: ByteBuf::with_capacity(1024),
46            logs: Logs::default(),
47            write_state: State::Start,
48            write_parent_state_stack: Vec::new(),
49            string_interner: StringInterner::new(),
50        }
51    }
52}
53
54impl Context {
55    #[cfg(not(target_family = "wasm"))]
56    fn new(input_bytes: Vec<u8>) -> Self {
57        Context {
58            input_bytes,
59            ..Default::default()
60        }
61    }
62
63    fn with<F, T>(f: F) -> T
64    where
65        F: FnOnce(&Context) -> T,
66    {
67        CONTEXT.with_borrow(f)
68    }
69
70    fn with_mut<F, T>(f: F) -> T
71    where
72        F: FnOnce(&mut Context) -> T,
73    {
74        CONTEXT.with_borrow_mut(f)
75    }
76}
77
78macro_rules! decorate_for_target {
79    ($(#[doc = $docs:tt])? fn $fn_name:ident($($args:tt)*) -> $ret:ty {
80        $($body:tt)*
81    }) => {
82        #[cfg(target_family = "wasm")]
83        $(#[doc = $docs])?
84        #[export_name = concat!("_", stringify!($fn_name))]
85        extern "C" fn $fn_name($($args)*) -> $ret {
86            $($body)*
87        }
88        #[cfg(not(target_family = "wasm"))]
89        $(#[doc = $docs])?
90        pub fn $fn_name($($args)*) -> $ret {
91            $($body)*
92        }
93    }
94}
95
96pub(crate) use decorate_for_target;
97
98use crate::log::Logs;
99
100#[cfg(target_family = "wasm")]
101#[export_name = "initialize"]
102extern "C" fn initialize(input_len: usize) -> *const u8 {
103    CONTEXT.with_borrow_mut(|context| {
104        *context = Context::default();
105        context.input_bytes = vec![0; input_len];
106        context.input_bytes.as_ptr()
107    })
108}
109
110#[cfg(not(target_family = "wasm"))]
111pub fn initialize_from_msgpack_bytes(bytes: Vec<u8>) {
112    CONTEXT.with_borrow_mut(|context| {
113        use std::mem;
114
115        let string_interner = mem::take(&mut context.string_interner);
116        *context = Context::new(bytes);
117        context.string_interner = string_interner;
118    })
119}
120
121#[cfg(target_family = "wasm")]
122#[export_name = "finalize"]
123extern "C" fn finalize() -> *const usize {
124    Context::with(|context| {
125        OUTPUT_AND_LOG_PTRS.with_borrow_mut(|output_and_log_ptrs| {
126            let output = context.output_bytes.as_vec();
127            output_and_log_ptrs[0] = output.as_ptr() as usize;
128            output_and_log_ptrs[1] = output.len();
129            let (log_offset1, log_len1, log_offset2, log_len2) = context.logs.read_ptrs();
130            output_and_log_ptrs[2] = log_offset1 as _;
131            output_and_log_ptrs[3] = log_len1;
132            output_and_log_ptrs[4] = log_offset2 as _;
133            output_and_log_ptrs[5] = log_len2;
134            output_and_log_ptrs.as_ptr()
135        })
136    })
137}
138
139decorate_for_target! {
140    fn shopify_function_intern_utf8_str(len: usize) -> DoubleUsize {
141        Context::with_mut(|context| {
142            let (id, ptr) = context.string_interner.preallocate(len);
143            ((id as DoubleUsize) << usize::BITS) | (ptr as DoubleUsize)
144        })
145    }
146}