forte-macros 0.4.0

//! # Integer Key Formatting for PK/SK
//!
//! When integer types are used as PK (Partition Key) or SK (Sort Key) fields,
//! they are zero-padded to their maximum decimal digit width so that
//! lexicographic (string) sort order matches numeric sort order.
//!
//! ## Unsigned types
//!
//! | Type          | Width | Example              |
//! |---------------|-------|----------------------|
//! | `u8`          | 3     | `042`                |
//! | `u16`         | 5     | `00042`              |
//! | `u32`         | 10    | `0000000042`         |
//! | `u64`/`usize` | 20   | `00000000000000000042` |
//!
//! ## Signed types (offset encoding)
//!
//! Signed integers are converted to unsigned by adding an offset equal to
//! `|T::MIN|` (i.e. `2^(bits-1)`), then zero-padded to the same width as the
//! corresponding unsigned type. This maps the full signed range onto `0..=U::MAX`
//! while preserving numeric order.
//!
//! | Type          | Offset           | Width | MIN → | 0 →  | MAX →  |
//! |---------------|------------------|-------|-------|------|--------|
//! | `i8`          | 128              | 3     | `000` | `128` | `255` |
//! | `i16`         | 32768            | 5     | `00000` | `32768` | `65535` |
//! | `i32`         | 2147483648       | 10    | `0000000000` | `2147483648` | `4294967295` |
//! | `i64`/`isize` | 9223372036854775808 | 20 | `00000000000000000000` | `09223372036854775808` | `18446744073709551615` |
//!
//! `usize` is always treated as `u64`, and `isize` as `i64`.

use proc_macro::TokenStream;
use quote::{format_ident, quote};
use syn::{parse_macro_input, Fields, ItemStruct};

fn format_placeholder(ty: &syn::Type) -> String {
    if let syn::Type::Path(type_path) = ty {
        if let Some(segment) = type_path.path.segments.last() {
            match segment.ident.to_string().as_str() {
                "u8" | "i8" => return "{:03}".to_string(),
                "u16" | "i16" => return "{:05}".to_string(),
                "u32" | "i32" => return "{:010}".to_string(),
                "u64" | "i64" | "usize" | "isize" => return "{:020}".to_string(),
                _ => {}
            }
        }
    }
    "{}".to_string()
}

fn wrap_expr(ty: &syn::Type, expr: proc_macro2::TokenStream) -> proc_macro2::TokenStream {
    if let syn::Type::Path(type_path) = ty {
        if let Some(segment) = type_path.path.segments.last() {
            match segment.ident.to_string().as_str() {
                "i8" => return quote! { (#expr as u8).wrapping_add(128u8) },
                "i16" => return quote! { (#expr as u16).wrapping_add(32768u16) },
                "i32" => return quote! { (#expr as u32).wrapping_add(2147483648u32) },
                "i64" | "isize" => return quote! { (#expr as u64).wrapping_add(9223372036854775808u64) },
                _ => {}
            }
        }
    }
    expr
}

#[proc_macro_attribute]
pub fn forte_doc(_attr: TokenStream, item: TokenStream) -> TokenStream {
    let input = parse_macro_input!(item as ItemStruct);

    let name = &input.ident;
    let vis = &input.vis;
    let get_name = format_ident!("{}Get", name);
    let put_name = format_ident!("{}Put", name);
    let query_name = format_ident!("{}Query", name);
    let delete_name = format_ident!("{}Delete", name);

    let fields = match &input.fields {
        Fields::Named(fields) => &fields.named,
        _ => panic!("forte_doc only supports named fields"),
    };

    let pk_fields: Vec<_> = fields
        .iter()
        .filter(|f| f.attrs.iter().any(|a| a.path().is_ident("pk")))
        .collect();

    let sk_fields: Vec<_> = fields
        .iter()
        .filter(|f| f.attrs.iter().any(|a| a.path().is_ident("sk")))
        .collect();

    let pk_field_names: Vec<_> = pk_fields.iter().map(|f| &f.ident).collect();
    let pk_field_types: Vec<_> = pk_fields.iter().map(|f| &f.ty).collect();
    let sk_field_names: Vec<_> = sk_fields.iter().map(|f| &f.ident).collect();
    let sk_field_types: Vec<_> = sk_fields.iter().map(|f| &f.ty).collect();

    let get_pk_fields: Vec<_> = pk_field_names
        .iter()
        .zip(pk_field_types.iter())
        .map(|(name, ty)| {
            let field_name = format_ident!("pk_{}", name.as_ref().unwrap());
            quote! { pub #field_name: #ty }
        })
        .collect();

    let get_sk_fields: Vec<_> = sk_field_names
        .iter()
        .zip(sk_field_types.iter())
        .map(|(name, ty)| {
            let field_name = format_ident!("sk_{}", name.as_ref().unwrap());
            quote! { pub #field_name: #ty }
        })
        .collect();

    let query_pk_fields: Vec<_> = pk_field_names
        .iter()
        .zip(pk_field_types.iter())
        .map(|(name, ty)| {
            let field_name = format_ident!("pk_{}", name.as_ref().unwrap());
            quote! { pub #field_name: #ty }
        })
        .collect();

    let query_sk_fields: Vec<_> = sk_field_names
        .iter()
        .zip(sk_field_types.iter())
        .map(|(name, ty)| {
            let field_name = format_ident!("sk_{}", name.as_ref().unwrap());
            quote! { pub #field_name: Option<#ty> }
        })
        .collect();

    let query_pk_str = if pk_fields.is_empty() {
        let name_str = name.to_string();
        quote! { #name_str.to_string() }
    } else {
        let name_str = name.to_string();
        let pk_format_parts: Vec<_> = pk_field_names
            .iter()
            .zip(pk_field_types.iter())
            .map(|(n, ty)| {
                let name_str = n.as_ref().unwrap().to_string();
                format!("{}={}", name_str, format_placeholder(ty))
            })
            .collect();
        let pk_format_string = format!("{}/{}", name_str, pk_format_parts.join("&"));
        let pk_format_args: Vec<_> = pk_field_names
            .iter()
            .zip(pk_field_types.iter())
            .map(|(n, ty)| {
                let field_name = format_ident!("pk_{}", n.as_ref().unwrap());
                wrap_expr(ty, quote! { self.#field_name })
            })
            .collect();
        quote! { format!(#pk_format_string, #(#pk_format_args),*) }
    };

    let query_sk_build: Vec<_> = sk_field_names
        .iter()
        .zip(sk_field_types.iter())
        .map(|(n, ty)| {
            let name_str = n.as_ref().unwrap().to_string();
            let field_name = format_ident!("sk_{}", n.as_ref().unwrap());
            let fmt = format!("{}={}", name_str, format_placeholder(ty));
            let val_expr = wrap_expr(ty, quote! { *v });
            quote! {
                if let Some(v) = &self.#field_name {
                    parts.push(format!(#fmt, #val_expr));
                } else {
                    break 'build;
                }
            }
        })
        .collect();

    let pk_str = if pk_fields.is_empty() {
        let name_str = name.to_string();
        quote! { #name_str.to_string() }
    } else {
        let name_str = name.to_string();
        let pk_format_parts: Vec<_> = pk_field_names
            .iter()
            .zip(pk_field_types.iter())
            .map(|(n, ty)| {
                let name_str = n.as_ref().unwrap().to_string();
                format!("{}={}", name_str, format_placeholder(ty))
            })
            .collect();
        let pk_format_string = format!("{}/{}", name_str, pk_format_parts.join("&"));
        let pk_format_args: Vec<_> = pk_field_names
            .iter()
            .zip(pk_field_types.iter())
            .map(|(n, ty)| {
                let field_name = format_ident!("pk_{}", n.as_ref().unwrap());
                wrap_expr(ty, quote! { self.#field_name })
            })
            .collect();
        quote! { format!(#pk_format_string, #(#pk_format_args),*) }
    };

    let sk_format_parts: Vec<_> = sk_field_names
        .iter()
        .zip(sk_field_types.iter())
        .map(|(n, ty)| {
            let name_str = n.as_ref().unwrap().to_string();
            format!("{}={}", name_str, format_placeholder(ty))
        })
        .collect();
    let sk_format_string = sk_format_parts.join("&");
    let sk_format_args: Vec<_> = sk_field_names
        .iter()
        .zip(sk_field_types.iter())
        .map(|(n, ty)| {
            let field_name = format_ident!("sk_{}", n.as_ref().unwrap());
            wrap_expr(ty, quote! { self.#field_name })
        })
        .collect();

    let put_pk_str = if pk_fields.is_empty() {
        let name_str = name.to_string();
        quote! { #name_str.to_string() }
    } else {
        let name_str = name.to_string();
        let pk_format_parts: Vec<_> = pk_field_names
            .iter()
            .zip(pk_field_types.iter())
            .map(|(n, ty)| {
                let name_str = n.as_ref().unwrap().to_string();
                format!("{}={}", name_str, format_placeholder(ty))
            })
            .collect();
        let pk_format_string = format!("{}/{}", name_str, pk_format_parts.join("&"));
        let pk_format_args: Vec<_> = pk_field_names
            .iter()
            .zip(pk_field_types.iter())
            .map(|(n, ty)| {
                let field_name = n.as_ref().unwrap();
                wrap_expr(ty, quote! { self.0.#field_name })
            })
            .collect();
        quote! { format!(#pk_format_string, #(#pk_format_args),*) }
    };

    let put_sk_format_args: Vec<_> = sk_field_names
        .iter()
        .zip(sk_field_types.iter())
        .map(|(n, ty)| {
            let field_name = n.as_ref().unwrap();
            wrap_expr(ty, quote! { self.0.#field_name })
        })
        .collect();

    let clean_fields: Vec<_> = fields
        .iter()
        .map(|f| {
            let mut f = f.clone();
            f.attrs.retain(|a| !a.path().is_ident("pk") && !a.path().is_ident("sk"));
            f
        })
        .collect();

    let expanded = quote! {
        #[derive(serde::Serialize, serde::Deserialize, Clone)]
        #vis struct #name {
            #(#clean_fields,)*
        }

        #vis struct #put_name(pub #name);

        impl forte_db::DbRequest for #put_name {
            type Output = ();
            fn prepare(self) -> forte_db::Prepared<Self::Output> {
                let pk = #put_pk_str;
                let sk = format!(#sk_format_string, #(#put_sk_format_args),*);
                let data = serde_json::to_vec(&self.0).expect("failed to serialize");
                forte_db::Prepared {
                    ops: vec![forte_db::DbOp::Put { pk, sk, data }],
                    parse: Box::new(|iter| {
                        match iter.next().ok_or_else(|| anyhow::anyhow!("missing result"))? {
                            forte_db::DbResult::Done => Ok(()),
                            _ => anyhow::bail!("unexpected result type"),
                        }
                    }),
                }
            }
        }

        #vis struct #get_name {
            #(#get_pk_fields,)*
            #(#get_sk_fields,)*
        }

        impl forte_db::DbRequest for #get_name {
            type Output = Option<#name>;
            fn prepare(self) -> forte_db::Prepared<Self::Output> {
                let pk = #pk_str;
                let sk = format!(#sk_format_string, #(#sk_format_args),*);
                forte_db::Prepared {
                    ops: vec![forte_db::DbOp::Get { pk, sk }],
                    parse: Box::new(|iter| {
                        match iter.next().ok_or_else(|| anyhow::anyhow!("missing result"))? {
                            forte_db::DbResult::Single(opt) => {
                                opt.map(|data| serde_json::from_slice(&data))
                                    .transpose()
                                    .map_err(Into::into)
                            }
                            _ => anyhow::bail!("unexpected result type"),
                        }
                    }),
                }
            }
        }

        #vis struct #query_name {
            #(#query_pk_fields,)*
            #(#query_sk_fields,)*
            pub limit: Option<usize>,
        }

        impl forte_db::DbRequest for #query_name {
            type Output = Vec<#name>;
            fn prepare(self) -> forte_db::Prepared<Self::Output> {
                let pk = #query_pk_str;
                let after_sk: Option<String> = {
                    let mut parts: Vec<String> = Vec::new();
                    'build: {
                        #(#query_sk_build)*
                    }
                    if parts.is_empty() { None } else { Some(parts.join("&")) }
                };
                let limit = self.limit;
                forte_db::Prepared {
                    ops: vec![forte_db::DbOp::Query { pk, after_sk, limit }],
                    parse: Box::new(|iter| {
                        match iter.next().ok_or_else(|| anyhow::anyhow!("missing result"))? {
                            forte_db::DbResult::Multiple(items) => {
                                items.into_iter()
                                    .map(|(_sk, data)| serde_json::from_slice(&data))
                                    .collect::<Result<Vec<_>, _>>()
                                    .map_err(Into::into)
                            }
                            _ => anyhow::bail!("unexpected result type"),
                        }
                    }),
                }
            }
        }

        #vis struct #delete_name {
            #(#get_pk_fields,)*
            #(#get_sk_fields,)*
        }

        impl forte_db::DbRequest for #delete_name {
            type Output = ();
            fn prepare(self) -> forte_db::Prepared<Self::Output> {
                let pk = #pk_str;
                let sk = format!(#sk_format_string, #(#sk_format_args),*);
                forte_db::Prepared {
                    ops: vec![forte_db::DbOp::Delete { pk, sk }],
                    parse: Box::new(|iter| {
                        match iter.next().ok_or_else(|| anyhow::anyhow!("missing result"))? {
                            forte_db::DbResult::Done => Ok(()),
                            _ => anyhow::bail!("unexpected result type"),
                        }
                    }),
                }
            }
        }
    };

    TokenStream::from(expanded)
}