// Copyright 2017 The Chromium OS Authors. All rights reserved.
// Use of this source code is governed by a BSD-style license that can be
// found in the LICENSE file.

//! Represents an address in the guest's memory space.

use std::cmp::{Eq, Ord, Ordering, PartialEq, PartialOrd};
use std::fmt::{self, Display};
use std::ops::{BitAnd, BitOr};

/// Represents an Address in the guest's memory.
#[derive(Clone, Copy, Debug)]
pub struct GuestAddress(pub u64);

impl GuestAddress {
    /// Returns the offset from this address to the given base address.
    ///
    /// # Examples
    ///
    /// ```
    /// # use sys_util::GuestAddress;
    ///   let base = GuestAddress(0x100);
    ///   let addr = GuestAddress(0x150);
    ///   assert_eq!(addr.offset_from(base), 0x50u64);
    /// ```
    pub fn offset_from(self, base: GuestAddress) -> u64 {
        self.0 - base.0
    }

    /// Returns the address as a u64 offset from 0x0.
    /// Use this when a raw number is needed to pass to the kernel.
    pub fn offset(self) -> u64 {
        self.0
    }

    /// Returns the result of the add or None if there is overflow.
    pub fn checked_add(self, other: u64) -> Option<GuestAddress> {
        self.0.checked_add(other).map(GuestAddress)
    }

    /// Returns the result of the base address + the size.
    /// Only use this when `offset` is guaranteed not to overflow.
    pub fn unchecked_add(self, offset: u64) -> GuestAddress {
        GuestAddress(self.0 + offset)
    }

    /// Returns the result of the subtraction of None if there is underflow.
    pub fn checked_sub(self, other: u64) -> Option<GuestAddress> {
        self.0.checked_sub(other).map(GuestAddress)
    }

    /// Returns the bitwise and of the address with the given mask.
    pub fn mask(self, mask: u64) -> GuestAddress {
        GuestAddress(self.0 & mask as u64)
    }
}

impl BitAnd<u64> for GuestAddress {
    type Output = GuestAddress;

    fn bitand(self, other: u64) -> GuestAddress {
        GuestAddress(self.0 & other as u64)
    }
}

impl BitOr<u64> for GuestAddress {
    type Output = GuestAddress;

    fn bitor(self, other: u64) -> GuestAddress {
        GuestAddress(self.0 | other as u64)
    }
}

impl PartialEq for GuestAddress {
    fn eq(&self, other: &GuestAddress) -> bool {
        self.0 == other.0
    }
}
impl Eq for GuestAddress {}

impl Ord for GuestAddress {
    fn cmp(&self, other: &GuestAddress) -> Ordering {
        self.0.cmp(&other.0)
    }
}

impl PartialOrd for GuestAddress {
    fn partial_cmp(&self, other: &GuestAddress) -> Option<Ordering> {
        Some(self.cmp(other))
    }
}

impl Display for GuestAddress {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        write!(f, "{:#x}", self.0)
    }
}

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

    #[test]
    fn equals() {
        let a = GuestAddress(0x300);
        let b = GuestAddress(0x300);
        let c = GuestAddress(0x301);
        assert_eq!(a, b);
        assert_eq!(b, a);
        assert_ne!(a, c);
        assert_ne!(c, a);
    }

    #[test]
    fn cmp() {
        let a = GuestAddress(0x300);
        let b = GuestAddress(0x301);
        assert!(a < b);
        assert!(b > a);
        assert!(!(a < a));
    }

    #[test]
    fn mask() {
        let a = GuestAddress(0x5050);
        assert_eq!(GuestAddress(0x5000), a & 0xff00u64);
        assert_eq!(GuestAddress(0x5055), a | 0x0005u64);
    }

    #[test]
    fn add_sub() {
        let a = GuestAddress(0x50);
        let b = GuestAddress(0x60);
        assert_eq!(Some(GuestAddress(0xb0)), a.checked_add(0x60));
        assert_eq!(0x10, b.offset_from(a));
    }

    #[test]
    fn checked_add_overflow() {
        let a = GuestAddress(0xffffffffffffff55);
        assert_eq!(Some(GuestAddress(0xffffffffffffff57)), a.checked_add(2));
        assert!(a.checked_add(0xf0).is_none());
    }
}