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, is_unsized_vec};
use crate::generators::persist::PersistGenerator;

impl PersistGenerator {
    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 version_fn = self.gen_table_version_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 pk_gen_state_fn = self.gen_table_pk_gen_state_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
                #version_fn
                #select_fn
                #select_range_fn
                #insert_fn
                #reinsert_fn
                #upsert_fn
                #count_fn
                #get_next_fn
                #pk_gen_state_fn
                #iter_with_fn
                #iter_with_async_fn
                #system_info_fn
                #vacuum_fn
            }
        }
    }

    fn gen_table_new_fn(&self) -> TokenStream {
        let name_generator = WorktableNameGenerator::from_table_name(self.name.to_string());
        let ident = name_generator.get_work_table_ident();
        let table_name = name_generator.get_work_table_literal_name();
        let task = name_generator.get_persistence_task_ident();
        let space_ident = name_generator.get_space_file_ident();
        let pk_type = name_generator.get_primary_key_type_ident();
        let const_name = name_generator.get_page_inner_size_const_ident();
        let secondary_index_events = name_generator.get_space_secondary_index_events_ident();
        let avt_index_ident = name_generator.get_available_indexes_ident();

        let pk_types = &self
            .columns
            .primary_keys
            .iter()
            .map(|i| {
                self.columns
                    .columns_map
                    .get(i)
                    .expect("should exist as got from definition")
                    .to_string()
            })
            .collect::<Vec<_>>();
        let pk_types_unsized = is_unsized_vec(pk_types);
        let index_setup = if pk_types_unsized {
            quote! {
                inner.primary_index = std::sync::Arc::new(PrimaryIndex {
                    pk_map: IndexMap::<#pk_type, OffsetEqLink<#const_name>, UnsizedNode<_>>::with_maximum_node_size(#const_name),
                    reverse_pk_map: IndexMap::new(),
                });
            }
        } else {
            quote! {
                let size = get_index_page_size_from_data_length::<#pk_type>(#const_name);
                inner.primary_index = std::sync::Arc::new(PrimaryIndex {
                    pk_map: IndexMap::<_, OffsetEqLink<#const_name>>::with_maximum_node_size(size),
                    reverse_pk_map: IndexMap::new(),
                });
            }
        };

        quote! {
            impl<E, C> PersistedWorkTable<E> for #ident
            where
                E: PersistenceEngine<
                    <<#pk_type as TablePrimaryKey>::Generator as PrimaryKeyGeneratorState>::State,
                    #pk_type,
                    #secondary_index_events,
                    #avt_index_ident,
                    Config=C
                > + Send
                    + 'static,
                C: Clone + PersistenceConfig,
            {
                async fn new(engine: E) -> eyre::Result<Self> {
                    let mut inner = WorkTable::default();
                    inner.table_name = #table_name;
                    #index_setup
                    core::result::Result::Ok(Self(
                        inner,
                        #task::run_engine(engine)
                    ))
                }

                async fn load(engine: E) -> eyre::Result<Self> {
                    let table_path = engine.config().table_path();
                    if !std::path::Path::new(table_path).exists() {
                        return Self::new(engine).await;
                    };
                    let space = #space_ident::parse_file(table_path).await?;
                    let table = space.into_worktable(engine).await;
                    Ok(table)
                }
            }
        }
    }

    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_version_fn(&self) -> TokenStream {
        let version = self.version;

        quote! {
            pub fn version() -> u32 {
                #version
            }
        }
    }

    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();
        let secondary_events_ident = name_generator.get_space_secondary_index_events_ident();

        quote! {
            pub fn insert(&self, row: #row_type) -> core::result::Result<#primary_key_type, WorkTableError> {
                let (op, res) = self.0.insert_cdc::<#secondary_events_ident>(row);
                if let Some(op) = op {
                    self.1.apply_operation(op);
                }
                res
            }
        }
    }

    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();
        let secondary_events_ident = name_generator.get_space_secondary_index_events_ident();

        quote! {
            pub async fn reinsert(&self, row_old: #row_type, row_new: #row_type) -> core::result::Result<#primary_key_type, WorkTableError> {
                let (op, res) = self.0.reinsert_cdc::<#secondary_events_ident>(row_old, row_new);
                if let Some(op) = op {
                    self.1.apply_operation(op);
                }
                res
            }
        }
    }

    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_pk_gen_state_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();

        quote! {
            pub fn pk_gen_state(&self) -> <<#primary_key_type as TablePrimaryKey>::Generator as PrimaryKeyGeneratorState>::State {
                self.0.pk_gen.get_state()
            }
        }
    }

    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 secondary_index_events = name_generator.get_space_secondary_index_events_ident();
        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,
                    _,
                    #secondary_index_events
                >::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),
                ).with_persistence(self.1.vacuum_sink()))
            }
        }
    }
}