Skip to main content

omega_cache_macros/
lib.rs

1use crate::ast::{AdmissionInput, BackoffInput, CacheInput, EngineInput, MetricsInput};
2use proc_macro::TokenStream;
3use quote::{ToTokens, quote};
4use syn::parse_macro_input;
5
6mod ast;
7mod parse;
8
9#[proc_macro]
10pub fn cache(input: TokenStream) -> TokenStream {
11    let input = parse_macro_input!(input as CacheInput);
12
13    let expanded = quote! {#input};
14
15    TokenStream::from(expanded)
16}
17
18impl ToTokens for CacheInput {
19    fn to_tokens(&self, tokens: &mut proc_macro2::TokenStream) {
20        let engine = &self.engine;
21        let admission_policy = &self.admission_policy;
22
23        tokens.extend(quote! {
24            crate::Cache::new(#engine, #admission_policy)
25        });
26    }
27}
28
29impl ToTokens for AdmissionInput {
30    fn to_tokens(&self, tokens: &mut proc_macro2::TokenStream) {
31        match self {
32            AdmissionInput::Always => tokens.extend(quote! { crate::AlwaysAdmission::new() }),
33            AdmissionInput::Frequent(frequent) => {
34                let frequent = frequent.as_ref();
35                let cms_width = &frequent.count_min_sketch.width;
36                let cms_height = &frequent.count_min_sketch.depth;
37                let decay_threshold = &frequent.decay_threshold;
38
39                let extend = quote! {
40                    crate::FrequentPolicy::new(#cms_width, #cms_height, #decay_threshold)
41                };
42
43                tokens.extend(extend)
44            }
45        }
46    }
47}
48
49impl ToTokens for EngineInput {
50    fn to_tokens(&self, tokens: &mut proc_macro2::TokenStream) {
51        match self {
52            EngineInput::Clock(clock) => {
53                let capacity = &clock.capacity;
54                let backoff = &clock.backoff;
55                let metrics = &clock.metrics;
56
57                tokens.extend(quote! {
58                    crate::clock::ClockCache::new(#capacity, #backoff, #metrics)
59                });
60            }
61            EngineInput::S3FIFO(s3fifo) => {
62                let capacity = &s3fifo.capacity;
63                let backoff = &s3fifo.backoff;
64                let metrics = &s3fifo.metrics;
65
66                tokens.extend(quote! {
67                    crate::s3fifo::S3FIFOCache::new(#capacity, #backoff, #metrics)
68                });
69            }
70        }
71    }
72}
73
74impl ToTokens for BackoffInput {
75    fn to_tokens(&self, tokens: &mut proc_macro2::TokenStream) {
76        let policy = &self.policy;
77        let limit = &self.limit;
78
79        let extend = quote! {
80            crate::core::backoff::BackoffConfig { policy: #policy, limit: #limit }
81        };
82
83        tokens.extend(extend);
84    }
85}
86
87impl ToTokens for MetricsInput {
88    fn to_tokens(&self, tokens: &mut proc_macro2::TokenStream) {
89        let shards = &self.shards;
90        let latency_samples = &self.latency_samples;
91
92        tokens.extend(quote! {
93            crate::metrics::MetricsConfig::new(#shards, #latency_samples)
94        });
95    }
96}