oar_ocr/core/

macros.rs

//! Macros for the OCR pipeline.
//!
//! This module provides utility macros to reduce code duplication across
//! the OCR pipeline, particularly for builder patterns and metrics collection.

/// Central task registry macro that defines all tasks in a single location.
///
/// This macro uses the "callback pattern" - it takes a callback macro name and
/// invokes it with the task registry data. Different consumers can process
/// the same data differently.
///
/// # Task Entry Format
///
/// Each task is defined as:
/// ```text
/// TaskName {
///     output: OutputType,           // fully qualified: $crate::domain::tasks::*
///     adapter: AdapterType,         // fully qualified: $crate::domain::adapters::*
///     input: image,
///     constructor: constructor_name,
///     conversion: into_method_name,
///     name: "snake_case_name",
///     doc: "Documentation string",
/// }
/// ```
///
/// # Benefits of Fully Qualified Paths
///
/// Using `$crate::` paths means consumer modules don't need to import
/// adapter/output types - the macro expansion includes the full paths.
#[macro_export]
macro_rules! with_task_registry {
    ($callback:path) => {
        $callback! {
            TextDetection {
                output: $crate::domain::tasks::TextDetectionOutput,
                adapter: $crate::domain::adapters::TextDetectionAdapter,
                input: image,
                constructor: text_detection,
                conversion: into_text_detection,
                name: "text_detection",
                doc: "Text detection - locating text regions in images",
            },
            TextRecognition {
                output: $crate::domain::tasks::TextRecognitionOutput,
                adapter: $crate::domain::adapters::TextRecognitionAdapter,
                input: image,
                constructor: text_recognition,
                conversion: into_text_recognition,
                name: "text_recognition",
                doc: "Text recognition - converting text regions to strings",
            },
            DocumentOrientation {
                output: $crate::domain::tasks::DocumentOrientationOutput,
                adapter: $crate::domain::adapters::DocumentOrientationAdapter,
                input: image,
                constructor: document_orientation,
                conversion: into_document_orientation,
                name: "document_orientation",
                doc: "Document orientation classification",
            },
            TextLineOrientation {
                output: $crate::domain::tasks::TextLineOrientationOutput,
                adapter: $crate::domain::adapters::TextLineOrientationAdapter,
                input: image,
                constructor: text_line_orientation,
                conversion: into_text_line_orientation,
                name: "text_line_orientation",
                doc: "Text line orientation classification",
            },
            DocumentRectification {
                output: $crate::domain::tasks::DocumentRectificationOutput,
                adapter: $crate::domain::adapters::UVDocRectifierAdapter,
                input: image,
                constructor: document_rectification,
                conversion: into_document_rectification,
                name: "document_rectification",
                doc: "Document rectification/unwarp",
            },
            LayoutDetection {
                output: $crate::domain::tasks::LayoutDetectionOutput,
                adapter: $crate::domain::adapters::LayoutDetectionAdapter,
                input: image,
                constructor: layout_detection,
                conversion: into_layout_detection,
                name: "layout_detection",
                doc: "Layout detection/analysis",
            },
            TableCellDetection {
                output: $crate::domain::tasks::TableCellDetectionOutput,
                adapter: $crate::domain::adapters::TableCellDetectionAdapter,
                input: image,
                constructor: table_cell_detection,
                conversion: into_table_cell_detection,
                name: "table_cell_detection",
                doc: "Table cell detection - locating cells within table regions",
            },
            FormulaRecognition {
                output: $crate::domain::tasks::FormulaRecognitionOutput,
                adapter: $crate::domain::adapters::FormulaRecognitionAdapter,
                input: image,
                constructor: formula_recognition,
                conversion: into_formula_recognition,
                name: "formula_recognition",
                doc: "Formula recognition - converting mathematical formulas to LaTeX",
            },
            SealTextDetection {
                output: $crate::domain::tasks::SealTextDetectionOutput,
                adapter: $crate::domain::adapters::SealTextDetectionAdapter,
                input: image,
                constructor: seal_text_detection,
                conversion: into_seal_text_detection,
                name: "seal_text_detection",
                doc: "Seal text detection - locating text regions in seal/stamp images",
            },
            TableClassification {
                output: $crate::domain::tasks::TableClassificationOutput,
                adapter: $crate::domain::adapters::TableClassificationAdapter,
                input: image,
                constructor: table_classification,
                conversion: into_table_classification,
                name: "table_classification",
                doc: "Table classification - classifying table images as wired or wireless",
            },
            TableStructureRecognition {
                output: $crate::domain::tasks::TableStructureRecognitionOutput,
                adapter: $crate::domain::adapters::TableStructureRecognitionAdapter,
                input: image,
                constructor: table_structure_recognition,
                conversion: into_table_structure_recognition,
                name: "table_structure_recognition",
                doc: "Table structure recognition - recognizing table structure as HTML with bboxes",
            }
        }
    };
}

/// Generates the TaskType enum from the task registry.
#[macro_export]
macro_rules! impl_task_type_enum {
    ($(
        $task:ident {
            output: $output:ty,
            adapter: $adapter:ty,
            input: $input_kind:ident,
            constructor: $constructor:ident,
            conversion: $conversion:ident,
            name: $name:literal,
            doc: $doc:literal,
        }
    ),* $(,)?) => {
        /// Represents the type of OCR task being performed.
        #[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, serde::Serialize, serde::Deserialize)]
        pub enum TaskType {
            $(
                #[doc = $doc]
                $task,
            )*
        }

        impl TaskType {
            /// Returns a human-readable name for the task type.
            pub fn name(&self) -> &'static str {
                match self {
                    $(TaskType::$task => $name,)*
                }
            }
        }
    };
}

/// Generates the DynTaskOutput enum and its methods from the task registry.
#[macro_export]
macro_rules! impl_dyn_task_output {
    ($(
        $task:ident {
            output: $output:ty,
            adapter: $adapter:ty,
            input: $input_kind:ident,
            constructor: $constructor:ident,
            conversion: $conversion:ident,
            name: $name:literal,
            doc: $doc:literal,
        }
    ),* $(,)?) => {
        /// Type-erased output from dynamic adapter execution.
        ///
        /// This enum wraps all possible task output types to enable dynamic dispatch.
        #[derive(Debug, Clone)]
        pub enum DynTaskOutput {
            $(
                #[doc = $doc]
                $task($output),
            )*
        }

        impl DynTaskOutput {
            $(
                #[doc = concat!("Extracts ", stringify!($output), " if this is a ", stringify!($task), " variant.")]
                pub fn $conversion(self) -> Result<$output, $crate::core::OCRError> {
                    match self {
                        Self::$task(output) => Ok(output),
                        _ => Err($crate::core::OCRError::InvalidInput {
                            message: format!(
                                concat!("Expected ", stringify!($task), " output, got {:?}"),
                                std::mem::discriminant(&self)
                            ),
                        }),
                    }
                }
            )*

            /// Returns the underlying task type for this output.
            pub fn task_type(&self) -> $crate::core::traits::task::TaskType {
                match self {
                    $(Self::$task(_) => $crate::core::traits::task::TaskType::$task,)*
                }
            }

            /// Creates an empty `DynTaskOutput` variant for the given task type.
            ///
            /// This is intended for registry wiring completeness checks and test scaffolding.
            pub fn empty_for(task_type: $crate::core::traits::task::TaskType) -> Self {
                match task_type {
                    $($crate::core::traits::task::TaskType::$task => {
                        Self::$task(<$output>::empty())
                    })*
                }
            }
        }
    };
}

/// Generates the TaskAdapter enum and its DynModelAdapter implementation.
///
/// All adapters are uniformly boxed to enable macro generation. The boxing overhead
/// is negligible since:
/// - Adapter creation happens once at build time
/// - Adapters are long-lived
/// - Inference time dominates any pointer overhead
/// - Most adapter internals are already indirect (Arc, Vec, etc.)
#[macro_export]
macro_rules! impl_task_adapter {
    ($(
        $task:ident {
            output: $output:ty,
            adapter: $adapter:ty,
            input: $input_kind:ident,
            constructor: $constructor:ident,
            conversion: $conversion:ident,
            name: $name:literal,
            doc: $doc:literal,
        }
    ),* $(,)?) => {
        /// Task-specific adapter enum with uniform Box storage.
        ///
        /// This enum directly holds concrete adapter types (boxed for uniform layout),
        /// avoiding the need for runtime downcast. Each variant corresponds to a
        /// specific task type and its adapter.
        ///
        /// # Benefits
        ///
        /// - **No runtime downcast**: Direct pattern matching on enum variants
        /// - **Compile-time exhaustiveness**: Adding a new task type requires handling it explicitly
        /// - **Type safety**: Each variant holds the exact adapter type for that task
        /// - **Uniform memory layout**: All variants are pointer-sized
        ///
        /// # Custom Adapters
        ///
        /// For testing or extension, use the `Custom` variant which wraps a `Box<dyn DynModelAdapter>`.
        #[derive(Debug)]
        pub enum TaskAdapter {
            $(
                #[doc = $doc]
                $task(Box<$adapter>),
            )*
            /// Custom adapter for testing or extension (wraps DynModelAdapter)
            Custom(Box<dyn DynModelAdapter>),
        }

        impl TaskAdapter {
            $(
                #[doc = concat!("Creates a TaskAdapter from a ", stringify!($adapter), ".")]
                pub fn $constructor(adapter: $adapter) -> Self {
                    Self::$task(Box::new(adapter))
                }
            )*

            /// Creates a TaskAdapter from a custom DynModelAdapter.
            pub fn custom<A: DynModelAdapter + 'static>(adapter: A) -> Self {
                Self::Custom(Box::new(adapter))
            }
        }

        impl DynModelAdapter for TaskAdapter {
            fn info(&self) -> AdapterInfo {
                match self {
                    $(Self::$task(a) => a.info(),)*
                    Self::Custom(a) => a.info(),
                }
            }

            fn task_type(&self) -> TaskType {
                match self {
                    $(Self::$task(_) => TaskType::$task,)*
                    Self::Custom(a) => a.task_type(),
                }
            }

            fn supports_batching(&self) -> bool {
                match self {
                    $(Self::$task(a) => a.supports_batching(),)*
                    Self::Custom(a) => a.supports_batching(),
                }
            }

            fn recommended_batch_size(&self) -> usize {
                match self {
                    $(Self::$task(a) => a.recommended_batch_size(),)*
                    Self::Custom(a) => a.recommended_batch_size(),
                }
            }

            fn execute_dyn(&self, input: DynTaskInput) -> Result<DynTaskOutput, OCRError> {
                match self {
                    $(
                        Self::$task(adapter) => {
                            $crate::impl_task_adapter!(@execute $input_kind, adapter, input, $task)
                        }
                    )*
                    Self::Custom(adapter) => adapter.execute_dyn(input),
                }
            }
        }
    };

    // Internal rule: execute for image input tasks
    (@execute image, $adapter:ident, $input:ident, $task:ident) => {{
        let image_input = match $input {
            DynTaskInput::Image(img) => img,
        };
        let output = $adapter.execute(image_input, None)?;
        Ok(DynTaskOutput::$task(output))
    }};
}

/// Macro to handle optional nested config initialization in builders.
///
/// This macro eliminates the repeated pattern of:
/// ```rust,no_run
/// // if self.config.field.is_none() {
/// //     self.config.field = Some(Type::new());
/// // }
/// ```
///
/// # Usage
///
/// ```rust,no_run
/// // Instead of:
/// // if self.config.orientation.is_none() {
/// //     self.config.orientation = Some(DocOrientationClassifierConfig::new());
/// // }
/// // if let Some(ref mut config) = self.config.orientation {
/// //     config.confidence_threshold = Some(threshold);
/// // }
///
/// // Use:
/// // with_nested!(self.config.orientation, DocOrientationClassifierConfig, config => {
/// //     config.confidence_threshold = Some(threshold);
/// // });
/// ```
#[macro_export]
macro_rules! with_nested {
    ($field:expr, $type:ty, $var:ident => $body:block) => {
        if $field.is_none() {
            $field = Some(<$type>::new());
        }
        if let Some(ref mut $var) = $field {
            $body
        }
    };
}

/// Macro to create pre-populated StageMetrics with common patterns.
///
/// This macro reduces duplication in metrics construction across stages.
///
/// # Usage
///
/// ```rust,no_run
/// // Instead of:
/// // StageMetrics::new(success_count, failure_count)
/// //     .with_processing_time(start_time.elapsed())
/// //     .with_info("stage", "cropping")
/// //     .with_info("batch_size", batch_size.to_string())
/// //     .with_info("parallel", parallel.to_string())
///
/// // Use:
/// // metrics!(success_count, failure_count, start_time; stage = "cropping", batch_size = batch_size, parallel = parallel)
/// // Or without timing:
/// // metrics!(success_count, failure_count; stage = "cropping", batch_size = batch_size)
/// ```
#[macro_export]
macro_rules! metrics {
    // With timing
    ($success:expr, $failure:expr, $start_time:expr; $($key:ident = $value:expr),*) => {
        {
            let mut metrics = $crate::pipeline::stages::StageMetrics::new($success, $failure);
            metrics = metrics.with_processing_time($start_time.elapsed());
            $(
                metrics = metrics.with_info(stringify!($key), $value.to_string());
            )*
            metrics
        }
    };
    // Without timing
    ($success:expr, $failure:expr; $($key:ident = $value:expr),*) => {
        {
            let mut metrics = $crate::pipeline::stages::StageMetrics::new($success, $failure);
            $(
                metrics = metrics.with_info(stringify!($key), $value.to_string());
            )*
            metrics
        }
    };
}

/// Comprehensive builder macro for generating common builder method patterns.
///
/// This macro generates multiple types of builder methods to reduce code duplication:
/// 1. Simple setters for direct field assignment
/// 2. Nested config setters using the `with_nested!` macro
/// 3. Enable/disable methods for optional features
/// 4. Dynamic batching configuration methods
///
/// # Usage
///
/// ```rust,no_run
/// // impl_complete_builder! {
/// //     builder: MyBuilder,
/// //     config_field: config,
///
/// //     // Simple setters
/// //     simple_setters: {
/// //         field_name: FieldType => "Documentation for the setter",
/// //     },
///
/// //     // Nested config setters
/// //     nested_setters: {
/// //         config_path: ConfigType => {
/// //             field_name: FieldType => "Documentation",
/// //         },
/// //     },
///
/// //     // Enable/disable methods
/// //     enable_methods: {
/// //         method_name => config_field: DefaultType => "Documentation",
/// //     },
/// // }
/// ```
#[macro_export]
macro_rules! impl_complete_builder {
    // Simple setters only
    (
        builder: $builder:ident,
        config_field: $config_field:ident,
        simple_setters: {
            $($simple_field:ident: $simple_type:ty => $simple_doc:literal),* $(,)?
        }
    ) => {
        impl $builder {
            $(
                #[doc = $simple_doc]
                pub fn $simple_field(mut self, value: $simple_type) -> Self {
                    self.$config_field.$simple_field = Some(value);
                    self
                }
            )*
        }
    };

    // Nested setters only
    (
        builder: $builder:ident,
        config_field: $config_field:ident,
        nested_setters: {
            $($nested_path:ident: $nested_type:ty => {
                $($nested_field:ident: $nested_field_type:ty => $nested_doc:literal),* $(,)?
            }),* $(,)?
        }
    ) => {
        impl $builder {
            $($(
                #[doc = $nested_doc]
                pub fn $nested_field(mut self, value: $nested_field_type) -> Self {
                    $crate::with_nested!(self.$config_field.$nested_path, $nested_type, config => {
                        config.$nested_field = Some(value);
                    });
                    self
                }
            )*)*
        }
    };

    // Enable methods only
    (
        builder: $builder:ident,
        config_field: $config_field:ident,
        enable_methods: {
            $($enable_method:ident => $enable_field:ident: $enable_type:ty => $enable_doc:literal),* $(,)?
        }
    ) => {
        impl $builder {
            $(
                #[doc = $enable_doc]
                pub fn $enable_method(mut self) -> Self {
                    self.$config_field.$enable_field = Some(<$enable_type>::default());
                    self
                }
            )*
        }
    };
}

/// Macro to implement `new()` and `with_common()` for config structs with per-module defaults.
#[macro_export]
macro_rules! impl_config_new_and_with_common {
    (
        $Config:ident,
        common_defaults: ($model_name_opt:expr, $batch_size_opt:expr),
        fields: { $( $field:ident : $default_expr:expr ),* $(,)? }
    ) => {
        impl $Config {
            /// Creates a new config instance with default values
            pub fn new() -> Self {
                Self {
                    common: $crate::core::config::builder::ModelInferenceConfig::with_defaults(
                        $model_name_opt, $batch_size_opt
                    ),
                    $( $field: $default_expr ),*
                }
            }
            /// Creates a new config instance using provided common configuration
            pub fn with_common(common: $crate::core::config::builder::ModelInferenceConfig) -> Self {
                Self {
                    common,
                    $( $field: $default_expr ),*
                }
            }
        }
    };
}

/// Macro to implement common builder methods for structs with a `ModelInferenceConfig` field.
#[macro_export]
macro_rules! impl_common_builder_methods {
    ($Builder:ident, $common_field:ident) => {
        impl $Builder {
            /// Sets the model path
            pub fn model_path(mut self, model_path: impl Into<std::path::PathBuf>) -> Self {
                self.$common_field = self.$common_field.model_path(model_path);
                self
            }
            /// Sets the model name
            pub fn model_name(mut self, model_name: impl Into<String>) -> Self {
                self.$common_field = self.$common_field.model_name(model_name);
                self
            }
            /// Sets the batch size
            pub fn batch_size(mut self, batch_size: usize) -> Self {
                self.$common_field = self.$common_field.batch_size(batch_size);
                self
            }
            /// Enables or disables logging
            pub fn enable_logging(mut self, enable: bool) -> Self {
                self.$common_field = self.$common_field.enable_logging(enable);
                self
            }
            /// Sets the ONNX Runtime session configuration
            pub fn ort_session(
                mut self,
                config: $crate::core::config::onnx::OrtSessionConfig,
            ) -> Self {
                self.$common_field = self.$common_field.ort_session(config);
                self
            }
        }
    };
}

/// Macro to inject common builder methods into an existing `impl Builder` block.
/// Use this inside `impl YourBuilder { ... }` and pass the field name that holds
/// `ModelInferenceConfig` (e.g., `common`).
#[macro_export]
macro_rules! common_builder_methods {
    ($common_field:ident) => {
        /// Sets the model path
        pub fn model_path(mut self, model_path: impl Into<std::path::PathBuf>) -> Self {
            self.$common_field = self.$common_field.model_path(model_path);
            self
        }
        /// Sets the model name
        pub fn model_name(mut self, model_name: impl Into<String>) -> Self {
            self.$common_field = self.$common_field.model_name(model_name);
            self
        }
        /// Sets the batch size
        pub fn batch_size(mut self, batch_size: usize) -> Self {
            self.$common_field = self.$common_field.batch_size(batch_size);
            self
        }
        /// Enables or disables logging
        pub fn enable_logging(mut self, enable: bool) -> Self {
            self.$common_field = self.$common_field.enable_logging(enable);
            self
        }
        /// Sets the ONNX Runtime session configuration
        pub fn ort_session(mut self, config: $crate::core::config::onnx::OrtSessionConfig) -> Self {
            self.$common_field = self.$common_field.ort_session(config);
            self
        }
    };
}

/// Macro to conditionally apply OrtSessionConfig to any builder that has `with_ort_config`.
///
/// This macro eliminates the repeated pattern:
/// ```rust,no_run
/// // let mut builder = SomeBuilder::new();
/// // if let Some(ort_config) = ort_config {
/// //     builder = builder.with_ort_config(ort_config);
/// // }
/// ```
///
/// Instead, use:
/// ```rust,no_run
/// // let builder = apply_ort_config!(SomeBuilder::new(), ort_config);
/// ```
///
/// # Usage
///
/// ```rust,no_run
/// // Works with any builder that has a `with_ort_config` method:
/// // let builder = apply_ort_config!(
/// //     DBModelBuilder::new()
/// //         .preprocess_config(config),
/// //     ort_config
/// // );
/// ```
#[macro_export]
macro_rules! apply_ort_config {
    ($builder:expr, $ort_config:expr) => {{
        let builder = $builder;
        if let Some(cfg) = $ort_config {
            builder.with_ort_config(cfg)
        } else {
            builder
        }
    }};
}

#[cfg(test)]
mod tests {

    // Test configuration structs
    #[derive(Debug, Default)]
    struct TestConfig {
        simple_field: Option<String>,
        nested_config: Option<NestedConfig>,
        enable_field: Option<EnabledFeature>,
    }

    #[derive(Debug, Default)]
    struct NestedConfig {
        nested_field: Option<i32>,
    }

    impl NestedConfig {
        fn new() -> Self {
            Self::default()
        }
    }

    #[derive(Debug, Default)]
    struct EnabledFeature {
        _enabled: bool,
    }

    // Test builder struct
    #[derive(Debug)]
    struct TestBuilder {
        config: TestConfig,
    }

    impl TestBuilder {
        fn new() -> Self {
            Self {
                config: TestConfig::default(),
            }
        }

        fn get_config(&self) -> &TestConfig {
            &self.config
        }
    }

    // Apply the macro to generate builder methods (separate calls for each type)
    impl_complete_builder! {
        builder: TestBuilder,
        config_field: config,
        simple_setters: {
            simple_field: String => "Sets a simple field value",
        }
    }

    impl_complete_builder! {
        builder: TestBuilder,
        config_field: config,
        nested_setters: {
            nested_config: NestedConfig => {
                nested_field: i32 => "Sets a nested field value",
            },
        }
    }

    impl_complete_builder! {
        builder: TestBuilder,
        config_field: config,
        enable_methods: {
            enable_feature => enable_field: EnabledFeature => "Enables a feature with default configuration",
        }
    }

    #[test]
    fn test_impl_complete_builder_nested_setter() {
        let builder = TestBuilder::new().nested_field(42);

        assert!(builder.get_config().nested_config.is_some());
        assert_eq!(
            builder
                .get_config()
                .nested_config
                .as_ref()
                .unwrap()
                .nested_field,
            Some(42)
        );
    }

    #[test]
    fn test_impl_complete_builder_enable_method() {
        let builder = TestBuilder::new().enable_feature();

        assert!(builder.get_config().enable_field.is_some());
    }

    #[test]
    fn test_impl_complete_builder_chaining() {
        let builder = TestBuilder::new()
            .simple_field("test".to_string())
            .nested_field(123)
            .enable_feature();

        let config = builder.get_config();
        assert_eq!(config.simple_field, Some("test".to_string()));
        assert!(config.nested_config.is_some());
        assert_eq!(
            config.nested_config.as_ref().unwrap().nested_field,
            Some(123)
        );
        assert!(config.enable_field.is_some());
    }
}