// Licensed to the Apache Software Foundation (ASF) under one

// or more contributor license agreements.  See the NOTICE file

// distributed with this work for additional information

// regarding copyright ownership.  The ASF licenses this file

// to you under the Apache License, Version 2.0 (the

// "License"); you may not use this file except in compliance

// with the License.  You may obtain a copy of the License at

//

//   http://www.apache.org/licenses/LICENSE-2.0

//

// Unless required by applicable law or agreed to in writing,

// software distributed under the License is distributed on an

// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY

// KIND, either express or implied.  See the License for the

// specific language governing permissions and limitations

// under the License.

use arrow_buffer::{bit_util, ArrowNativeType, Buffer, MutableBuffer};

/// Creates two [`Buffer`]s from an iterator of `Option`.

/// The first buffer corresponds to a bitmap buffer, the second one

/// corresponds to a values buffer.

/// # Safety

/// The caller must ensure that `iterator` is `TrustedLen`.

#[inline]

pub(crate) unsafe fn trusted_len_unzip<I, P, T>(iterator: I) -> (Buffer, Buffer)

where

    T: ArrowNativeType,

    P: std::borrow::Borrow<Option<T>>,

    I: Iterator<Item = P>,

{

    let (_, upper) = iterator.size_hint();

    let upper = upper.expect("trusted_len_unzip requires an upper limit");

    let len = upper * std::mem::size_of::<T>();

    let mut null = MutableBuffer::from_len_zeroed(upper.saturating_add(7) / 8);

    let mut buffer = MutableBuffer::new(len);

    let dst_null = null.as_mut_ptr();

    let mut dst = buffer.as_mut_ptr() as *mut T;

    for (i, item) in iterator.enumerate() {

        let item = item.borrow();

        if let Some(item) = item {

            std::ptr::write(dst, *item);

            bit_util::set_bit_raw(dst_null, i);

        } else {

            std::ptr::write(dst, T::default());

        }

        dst = dst.add(1);

    }

    assert_eq!(

        dst.offset_from(buffer.as_ptr() as *mut T) as usize,

        upper,

        "Trusted iterator length was not accurately reported"

    );

    buffer.set_len(len);

    (null.into(), buffer.into())

}

#[cfg(test)]

mod tests {

    use super::*;

    #[test]

    fn trusted_len_unzip_good() {

        let vec = [Some(1u32), None];

        let (null, buffer) = unsafe { trusted_len_unzip(vec.iter()) };

        assert_eq!(null.as_slice(), &[0b00000001]);

        assert_eq!(buffer.as_slice(), &[1u8, 0, 0, 0, 0, 0, 0, 0]);

    }

    #[test]

    #[should_panic(expected = "trusted_len_unzip requires an upper limit")]

    fn trusted_len_unzip_panic() {

        let iter = std::iter::repeat(Some(4i32));

        unsafe { trusted_len_unzip(iter) };

    }

}