worktable_macros 0.9.0

use convert_case::{Case, Casing};
use proc_macro2::{Ident, Span, TokenStream};
use quote::quote;

use crate::common::model::GeneratorType;
use crate::common::name_generator::WorktableNameGenerator;
use crate::generators::in_memory::InMemoryGenerator;

impl InMemoryGenerator {
    pub fn gen_table_impl(&self) -> TokenStream {
        let name_generator = WorktableNameGenerator::from_table_name(self.name.to_string());
        let ident = name_generator.get_work_table_ident();

        let persisted_impl = self.gen_table_new_fn();
        let name_fn = self.gen_table_name_fn();
        let select_fn = self.gen_table_select_fn();
        let select_range_fn = self.gen_table_select_range_fn();
        let insert_fn = self.gen_table_insert_fn();
        let reinsert_fn = self.gen_table_reinsert_fn();
        let upsert_fn = self.gen_table_upsert_fn();
        let get_next_fn = self.gen_table_get_next_fn();
        let iter_with_fn = self.gen_table_iter_with_fn();
        let iter_with_async_fn = self.gen_table_iter_with_async_fn();
        let count_fn = self.gen_table_count_fn();
        let system_info_fn = self.gen_system_info_fn();
        let vacuum_fn = self.gen_table_vacuum_fn();

        quote! {
            #persisted_impl
            impl #ident {
                #name_fn
                #select_fn
                #select_range_fn
                #insert_fn
                #reinsert_fn
                #upsert_fn
                #count_fn
                #get_next_fn
                #iter_with_fn
                #iter_with_async_fn
                #system_info_fn
                #vacuum_fn
            }
        }
    }

    fn gen_table_new_fn(&self) -> TokenStream {
        // InMemory tables don't have PersistedWorkTable impl
        quote! {}
    }

    fn gen_table_name_fn(&self) -> TokenStream {
        let name_generator = WorktableNameGenerator::from_table_name(self.name.to_string());
        let dir_name = name_generator.get_dir_name();

        quote! {
            pub fn name(&self) -> &'static str {
                &self.0.table_name
            }

            pub fn name_snake_case() -> &'static str {
                #dir_name
            }
        }
    }

    fn gen_table_select_fn(&self) -> TokenStream {
        let name_generator = WorktableNameGenerator::from_table_name(self.name.to_string());
        let row_type = name_generator.get_row_type_ident();
        let primary_key_type = name_generator.get_primary_key_type_ident();

        quote! {
            pub fn select<Pk>(&self, pk: Pk) -> Option<#row_type>
            where #primary_key_type: From<Pk> {
                self.0.select(pk.into())
            }
        }
    }

    fn gen_table_select_range_fn(&self) -> TokenStream {
        let name_generator = WorktableNameGenerator::from_table_name(self.name.to_string());
        let row_type = name_generator.get_row_type_ident();
        let primary_key_type = name_generator.get_primary_key_type_ident();
        let column_range_type = name_generator.get_column_range_type_ident();
        let row_fields_ident = name_generator.get_row_fields_enum_ident();

        let pk_sorted_by = if self.columns.primary_keys.len() == 1 {
            let pk_field = &self.columns.primary_keys[0];
            let pk_pascal = Ident::new(&pk_field.to_string().to_case(Case::Pascal), Span::mixed_site());
            quote! {
                SelectQueryBuilder::new_sorted(rows, #row_fields_ident::#pk_pascal)
            }
        } else {
            quote! {
                SelectQueryBuilder::new(rows)
            }
        };

        quote! {
            pub fn select_by_pk_range<R, Pk>(&self, range: R) -> SelectQueryBuilder<#row_type,
                                                                     impl DoubleEndedIterator<Item = #row_type> + '_,
                                                                     #column_range_type,
                                                                     #row_fields_ident>
            where
                #primary_key_type: From<Pk>,
                R: std::ops::RangeBounds<Pk>,
                Pk: Clone,
            {
                let converted_range = (
                    range.start_bound().map(|v| #primary_key_type::from(v.clone())),
                    range.end_bound().map(|v| #primary_key_type::from(v.clone())),
                );
                let rows = self.0.primary_index.pk_map
                    .range(converted_range)
                    .filter_map(|(_, link)| self.0.data.select_non_ghosted(link.0).ok());

                #pk_sorted_by
            }
        }
    }

    fn gen_table_insert_fn(&self) -> TokenStream {
        let name_generator = WorktableNameGenerator::from_table_name(self.name.to_string());
        let row_type = name_generator.get_row_type_ident();
        let primary_key_type = name_generator.get_primary_key_type_ident();

        quote! {
            pub fn insert(&self, row: #row_type) -> core::result::Result<#primary_key_type, WorkTableError> {
                self.0.insert(row)
            }
        }
    }

    fn gen_table_reinsert_fn(&self) -> TokenStream {
        let name_generator = WorktableNameGenerator::from_table_name(self.name.to_string());
        let row_type = name_generator.get_row_type_ident();
        let primary_key_type = name_generator.get_primary_key_type_ident();

        quote! {
            pub async fn reinsert(&self, row_old: #row_type, row_new: #row_type) -> core::result::Result<#primary_key_type, WorkTableError> {
                self.0.reinsert(row_old, row_new).await
            }
        }
    }

    fn gen_table_upsert_fn(&self) -> TokenStream {
        let name_generator = WorktableNameGenerator::from_table_name(self.name.to_string());
        let row_type = name_generator.get_row_type_ident();

        quote! {
            /// Inserts the row if its primary key is absent, updates it
            /// otherwise.
            ///
            /// Concurrency: **system-wide lock-free, not wait-free per call.**
            /// A retry is only taken when a concurrent delete or insert
            /// flipped this key's existence between the existence check and
            /// the operation, so some operation on this key completes on
            /// every iteration -- but under sustained adversarial churn on
            /// the same key an individual call can retry indefinitely. There
            /// is deliberately no retry limit: upsert is semantically
            /// infallible for primary-key conflicts, and a limit would trade
            /// theoretical starvation for real spurious errors. Each
            /// conflicting round yields to the scheduler before retrying so
            /// the interfering task can complete.
            pub async fn upsert(&self, row: #row_type) -> core::result::Result<(), WorkTableError> {
                let pk = row.get_primary_key();
                loop {
                    let need_to_update = self.0.primary_index.pk_map.get(&pk).is_some();
                    if need_to_update {
                        match self.update(row.clone()).await {
                            core::result::Result::Ok(_) => return core::result::Result::Ok(()),
                            // Row was deleted concurrently between the check and the
                            // update; retry as an insert.
                            core::result::Result::Err(WorkTableError::NotFound) => {
                                tokio::task::yield_now().await;
                                continue;
                            }
                            // Row is mid-flight: a concurrent insert publishes
                            // the primary-key entry before unghosting the row
                            // data (and insert takes no row lock), and a
                            // concurrent delete ghosts data it is about to
                            // unindex. Both are transient; retry.
                            core::result::Result::Err(WorkTableError::PagesError(e)) if e.is_row_absent() => {
                                tokio::task::yield_now().await;
                                continue;
                            }
                            core::result::Result::Err(e) => return core::result::Result::Err(e),
                        }
                    } else {
                        match self.insert(row.clone()) {
                            core::result::Result::Ok(_) => return core::result::Result::Ok(()),
                            // Row was inserted concurrently between the check and the
                            // insert; retry as an update. Secondary-index conflicts are
                            // real errors and are propagated. Progress is lock-free,
                            // not wait-free: a retry is only taken when a concurrent
                            // delete/insert flipped this key's existence between the
                            // check and the operation, so the system as a whole makes
                            // progress on every retry, but this call can in principle
                            // retry unboundedly under sustained same-key churn.
                            core::result::Result::Err(WorkTableError::PrimaryAlreadyExists) => {
                                tokio::task::yield_now().await;
                                continue;
                            }
                            core::result::Result::Err(e) => return core::result::Result::Err(e),
                        }
                    }
                }
            }
        }
    }

    fn gen_table_get_next_fn(&self) -> TokenStream {
        let name_generator = WorktableNameGenerator::from_table_name(self.name.to_string());
        let primary_key_type = name_generator.get_primary_key_type_ident();

        match self.columns.generator_type {
            GeneratorType::Custom | GeneratorType::Autoincrement => {
                quote! {
                    pub fn get_next_pk(&self) -> #primary_key_type {
                        self.0.get_next_pk()
                    }
                }
            }
            GeneratorType::None => {
                quote! {}
            }
        }
    }

    fn gen_table_iter_with_fn(&self) -> TokenStream {
        let name_generator = WorktableNameGenerator::from_table_name(self.name.to_string());
        let row_type = name_generator.get_row_type_ident();
        let inner = self.gen_table_iter_inner(quote! {
            f(data)?;
        });

        quote! {
            pub fn iter_with<
                F: Fn(#row_type) -> core::result::Result<(), WorkTableError>
            >(&self, f: F) -> core::result::Result<(), WorkTableError> {
                #inner
            }
        }
    }

    fn gen_table_iter_with_async_fn(&self) -> TokenStream {
        let name_generator = WorktableNameGenerator::from_table_name(self.name.to_string());
        let row_type = name_generator.get_row_type_ident();
        let inner = self.gen_table_iter_inner(quote! {
             f(data).await?;
        });

        quote! {
            pub async fn iter_with_async<
                F: Fn(#row_type) -> Fut,
                Fut: std::future::Future<Output = core::result::Result<(), WorkTableError>>
            >(&self, f: F) -> core::result::Result<(), WorkTableError> {
                #inner
            }
        }
    }

    fn gen_table_iter_inner(&self, func: TokenStream) -> TokenStream {
        quote! {
            let first = self.0.primary_index.pk_map.iter().next().map(|(k, v)| (k.clone(), v.0));
            let Some((mut k, link)) = first else {
                return Ok(())
            };

            let data = self.0.data.select_non_ghosted(link).map_err(WorkTableError::PagesError)?;
            #func

            let mut ind = false;
            while !ind {
                let next = {
                    let mut iter = self.0.primary_index.pk_map.range(k.clone()..);
                    let next = iter.next().map(|(k, v)| (k.clone(), v.0)).filter(|(key, _)| key != &k);
                    if next.is_some() {
                        next
                    } else {
                        iter.next().map(|(k, v)| (k.clone(), v.0))
                    }
                };
                if let Some((key, link)) = next {
                    let data = self.0.data.select_non_ghosted(link).map_err(WorkTableError::PagesError)?;
                   #func
                    k = key
                } else {
                    ind = true;
                };
            }

            core::result::Result::Ok(())
        }
    }

    fn gen_table_count_fn(&self) -> TokenStream {
        quote! {
            pub fn count(&self) -> usize {
                let count = self.0.primary_index.pk_map.len();
                count
            }
        }
    }

    fn gen_system_info_fn(&self) -> TokenStream {
        quote! {
            pub fn system_info(&self) -> SystemInfo {
                self.0.system_info()
            }
        }
    }

    fn gen_table_vacuum_fn(&self) -> TokenStream {
        let name_generator = WorktableNameGenerator::from_table_name(self.name.to_string());
        let table_name = name_generator.get_work_table_literal_name();
        let lock_type = name_generator.get_lock_type_ident();

        quote! {
            pub fn vacuum(&self) -> std::sync::Arc<dyn WorkTableVacuum + std::marker::Send + Sync> {
                std::sync::Arc::new(EmptyDataVacuum::<
                    _,
                    _,
                    _,
                    _,
                    _,
                    _,
                    #lock_type,
                    _
                >::new(
                    #table_name,
                    std::sync::Arc::clone(&self.0.data),
                    std::sync::Arc::clone(&self.0.lock_manager),
                    std::sync::Arc::clone(&self.0.primary_index),
                    std::sync::Arc::clone(&self.0.indexes),
                ))
            }
        }
    }
}