diskann-disk 0.55.0

/*
 * Copyright (c) Microsoft Corporation.
 * Licensed under the MIT license.
 */

use std::time::Duration;

// ContinuationGrant struct used for ochestrating chunkable operations.
pub enum ContinuationGrant {
    Continue,
    Yield(Duration),
    Stop,
}

// Contituation grant trait to be implemented by the client.
// This trait is used to get a continuation grant during chunk intervals while doing chunkable operations.
pub trait ContinuationTrackerTrait: Send + Sync + ContinuationCheckerTraitClone {
    fn get_continuation_grant(&self) -> ContinuationGrant;
}

// Naive implementation of the ContinuationGrantProviderTrait
// This implementation always returns ContinuationGrant::Continue.
#[derive(Default, Clone)]
pub struct NaiveContinuationTracker {}

impl ContinuationTrackerTrait for NaiveContinuationTracker {
    fn get_continuation_grant(&self) -> ContinuationGrant {
        ContinuationGrant::Continue
    }
}

pub trait ContinuationCheckerTraitClone {
    fn clone_box(&self) -> Box<dyn ContinuationTrackerTrait>;
}

impl<T> ContinuationCheckerTraitClone for T
where
    T: 'static + ContinuationTrackerTrait + Clone,
{
    fn clone_box(&self) -> Box<dyn ContinuationTrackerTrait> {
        Box::new(self.clone())
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn test_naive_continuation_tracker_default() {
        let tracker = NaiveContinuationTracker::default();
        // Verify it always returns Continue
        match tracker.get_continuation_grant() {
            ContinuationGrant::Continue => {}
            _ => panic!("Expected Continue"),
        }
    }

    #[test]
    fn test_naive_continuation_tracker_clone_box() {
        let tracker = NaiveContinuationTracker::default();
        let boxed = tracker.clone_box();

        // The boxed version should also return Continue
        match boxed.get_continuation_grant() {
            ContinuationGrant::Continue => {}
            _ => panic!("Expected Continue"),
        }
    }
}