// Copyright 2022 Google LLC
//
// Licensed 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
//
//     https://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 crate::block_device::*;
use crate::errors::*;
use crate::utils::*;

pub struct MemmapDevice {
    pub geom: BlockDeviceGeometry,
    pub flash_base: *mut u8,
}

impl MemmapDevice {
    pub fn new(geom: BlockDeviceGeometry, flash_base: *mut u8) -> Self {
        let result = MemmapDevice {
            geom: geom,
            flash_base: flash_base,
        };
        return result;
    }
}

impl MemmapDevice {
    pub fn check_is_block_sized(&self, data: &[u8]) -> Result<(), BFSErr> {
        dcheck!(data.len() == self.geom.block_size, BFSErr::OutOfBounds);
        return Ok(());
    }

    pub fn check_block_index(&self, iblock: usize) -> Result<(), BFSErr> {
        dcheck!(iblock < self.geom.block_count, BFSErr::OutOfBounds);
        return Ok(());
    }

    pub fn check_fits_in_block(&self, addr: usize, data: &[u8]) -> Result<(), BFSErr> {
        let block_a = addr / self.geom.block_size;
        let block_b = (addr + data.len() - 1) / self.geom.block_size;
        dcheck!(block_a == block_b, BFSErr::OutOfBounds);
        return Ok(());
    }

    pub fn check_read_block(&self, iblock: usize, block: &[u8]) -> Result<(), BFSErr> {
        self.check_block_index(iblock)?;
        self.check_fits_in_block(iblock * self.geom.block_size, block)?;
        return Ok(());
    }

    pub fn check_write_block(&self, iblock: usize, block: &[u8]) -> Result<(), BFSErr> {
        self.check_block_index(iblock)?;
        self.check_fits_in_block(iblock * self.geom.block_size, block)?;
        return Ok(());
    }

    pub fn check_erase_block(&self, iblock: usize) -> Result<(), BFSErr> {
        self.check_block_index(iblock)?;
        return Ok(());
    }

    pub fn check_erase_dirty_block(&self, iblock: usize) -> Result<(), BFSErr> {
        self.check_erase_block(iblock)?;
        return Ok(());
    }

    pub fn check_range(&self, addr: usize, size: usize) -> Result<(), BFSErr> {
        let bs = self.geom.block_size;
        let bc = self.geom.block_count;
        dcheck!(addr + size < (bc * bs), BFSErr::OutOfBounds);
        return Ok(());
    }

    pub fn check_read_range(&self, addr: usize, data: &mut [u8]) -> Result<(), BFSErr> {
        self.check_range(addr, data.len())?;
        return Ok(());
    }

    pub fn check_write_range(&self, addr: usize, data: &[u8]) -> Result<(), BFSErr> {
        self.check_range(addr, data.len())?;
        return Ok(());
    }

    /// Check that this overwrite only changes bits from 1->0
    pub fn check_overwrite_range(&self, addr: usize, data: &[u8]) -> Result<(), BFSErr> {
        self.check_range(addr, data.len())?;
        for i in 0..data.len() {
            let src = data[i];
            unsafe {
                let dst = self.flash_base.add(addr + i).read();
                dcheck!((src & !dst) == 0, BFSErr::BadOverwrite);
            }
        }
        return Ok(());
    }
}

impl BlockDevice for MemmapDevice {
    fn geom(&self) -> BlockDeviceGeometry {
        self.geom
    }

    /// Read a chunk of flash contained in a single block.
    fn read_block(&self, iblock: usize, block: &mut [u8]) -> Result<(), BFSErr> {
        self.check_read_block(iblock, block)?;

        let bs = self.geom.block_size;
        unsafe {
            self.flash_base
                .add(iblock * bs)
                .copy_to(block.as_mut_ptr(), block.len());
        }

        return Ok(());
    }

    /// Write a a chunk of flash contained in a single block.
    fn write_block(&mut self, iblock: usize, block: &[u8]) -> Result<(), BFSErr> {
        self.check_write_block(iblock, block)?;

        let bs = self.geom.block_size;
        unsafe {
            self.flash_base
                .add(iblock * bs)
                .copy_from(block.as_ptr(), bs);
        }

        return Ok(());
    }

    /// Erase an entire block of flash.
    fn erase_block(&mut self, iblock: usize) -> Result<(), BFSErr> {
        self.check_erase_block(iblock)?;

        let bs = self.geom.block_size;
        unsafe {
            self.flash_base.add(iblock * bs).write_bytes(0xFF, bs);
        }

        return Ok(());
    }

    /// Read a range of bytes in flash, checking first that those bytes have
    /// been written since they were erased.
    fn read_range(&self, addr: usize, data: &mut [u8]) -> Result<(), BFSErr> {
        self.check_read_range(addr, data)?;

        unsafe {
            self.flash_base
                .add(addr)
                .copy_to(data.as_mut_ptr(), data.len());
        }

        return Ok(());
    }

    /// Write to a range of bytes in flash, checking first that those bytes have
    /// not been written since they were erased.
    fn write_range(&mut self, addr: usize, data: &[u8]) -> Result<(), BFSErr> {
        self.check_write_range(addr, data)?;

        unsafe {
            self.flash_base
                .add(addr)
                .copy_from(data.as_ptr(), data.len());
        }

        return Ok(());
    }

    /// Write to a range of bytes in flash, but _do_ allow writing over dirty
    /// bytes - the result in flash will be the logical AND of the old and new
    /// data.
    fn overwrite_range(&mut self, addr: usize, data: &[u8]) -> Result<(), BFSErr> {
        self.check_overwrite_range(addr, data)?;

        unsafe {
            for i in 0..data.len() {
                let ptr = self.flash_base.add(addr + i);
                ptr.write(ptr.read() & data[i]);
            }
        }

        return Ok(());
    }
}

/// Overwriting existing data should write the logical AND of the old and new
/// data.

#[test]
fn test_overwrite() {
    let geom = BlockDeviceGeometry {
        block_size: 16,
        block_count: 4,
    };
    let mut buf: Vec<u8> = vec![0; geom.block_count * geom.block_size];
    let bd: &mut dyn BlockDevice = &mut MemmapDevice::new(geom, buf.as_mut_ptr());

    let write1: u32 = 0xFF0FF0FF;
    let write2: u32 = 0x0F00F000;

    assert_ok!(bd.write_range(0, &write1.to_le_bytes()));
    assert_ok!(bd.overwrite_range(0, &write2.to_le_bytes()));

    let mut result: u32 = 0;
    assert_ok!(bd.read_range(0, as_unsafe_blob_mut(&mut result)));
    assert_eq!(result, 0xFF0FF0FF & 0x0FF0FF00);
}

/// Trying to set bits that have already been cleared by a previous write should
/// cause a panic in the test device.

#[test]
#[should_panic]
fn test_bad_overwrite() {
    let geom = BlockDeviceGeometry {
        block_size: 16,
        block_count: 4,
    };
    let mut buf: Vec<u8> = vec![0; geom.block_count * geom.block_size];
    let bd: &mut dyn BlockDevice = &mut MemmapDevice::new(geom, buf.as_mut_ptr());

    let all_1: u32 = 0xFFFFFFFF;
    let all_0: u32 = 0x00000000;

    assert_ok!(bd.write_range(0, &all_0.to_le_bytes()));
    assert_err!(bd.overwrite_range(0, &all_1.to_le_bytes()));
}

/// Trying to read blocks that have never been written should cause a panic in
/// the test device.

#[test]
#[should_panic]
fn test_read_unwritten_block() {
    let geom = BlockDeviceGeometry {
        block_size: 16,
        block_count: 4,
    };
    let mut buf: Vec<u8> = vec![0; geom.block_count * geom.block_size];
    let bd: &mut dyn BlockDevice = &mut MemmapDevice::new(geom, buf.as_mut_ptr());

    let mut dst_block = vec![0; geom.block_size];
    assert_err!(bd.read_block(0, &mut dst_block));
}

/// Trying to read blocks that have been written and then erased should cause a
/// panic in the test device.

#[test]
#[should_panic]
fn test_read_erased_block() {
    let geom = BlockDeviceGeometry {
        block_size: 16,
        block_count: 4,
    };
    let mut buf: Vec<u8> = vec![0; geom.block_count * geom.block_size];
    let bd: &mut dyn BlockDevice = &mut MemmapDevice::new(geom, buf.as_mut_ptr());

    let mut src_block = vec![0; geom.block_size];
    let mut dst_block = vec![0; geom.block_size];

    for i in 0..geom.block_size {
        src_block[i] = i as u8;
    }

    assert_ok!(bd.write_block(0, &mut src_block));
    assert_ok!(bd.erase_block(0));
    assert_err!(bd.read_block(0, &mut dst_block));
}

/// All blocks in a device should be readable and writable.

#[test]
fn test_read_write() {
    let geom = BlockDeviceGeometry {
        block_size: 16,
        block_count: 4,
    };
    let mut buf: Vec<u8> = vec![0; geom.block_count * geom.block_size];
    let bd: &mut dyn BlockDevice = &mut MemmapDevice::new(geom, buf.as_mut_ptr());

    let mut src_block = vec![0; geom.block_size];
    let mut dst_block = vec![0; geom.block_size];

    for i in 0..geom.block_size {
        src_block[i] = i as u8;
    }

    // Writing clean blocks should succeed.
    for i in 0..geom.block_count {
        assert_ok!(bd.write_block(i, &src_block));
    }

    // Reading those blocks back should succeed, and the contents should match
    // the original block.
    for i in 0..geom.block_count {
        assert_ok!(bd.read_block(i, &mut dst_block));
        for j in 0..geom.block_size {
            assert_eq!(src_block[j], dst_block[j]);
        }
    }

    // Erasing blocks should succeed.
    for i in 0..geom.block_count {
        assert_ok!(bd.erase_block(i));
    }
}