// Panopticon - A libre program analysis library for machine code
// Copyright (C) 2014-2018  The Panopticon Developers
//
// This library is free software; you can redistribute it and/or
// modify it under the terms of the GNU Lesser General Public
// License as published by the Free Software Foundation; either
// version 2.1 of the License, or (at your option) any later version.
//
// This library is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU
// Lesser General Public License for more details.
//
// You should have received a copy of the GNU Lesser General Public
// License along with this library; if not, write to the Free Software
// Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston, MA  02110-1301  USA

//! Constraint on a variable.
//!
//! This type is used to constraint the values of a variable to a range of values. Constraints can
//! be a signed or unsigned range of `n`-bit integer.
//!
//! Used for narrowing in Abstract Interpretation.

use std::{u64,i64,cmp};
use num_traits::Saturating;
use quickcheck::{Arbitrary,Gen};
use {Value,Constant,Result};

/// A range of `n`-bit integers.
#[derive(Debug,Clone,PartialEq,Eq)]
pub enum Constraint {
    /// The empty range containing no values.
    Empty{
        /// Bitsize of the integer this range constraints.
        bits: usize,
    },

    /// A range `from`-`to`, both inclusive.
    Unsigned{
        /// Lower bound (inclusive).
        from: u64,
        /// Upper bound (inclusive).
        to: u64,
        /// Bitsize of the integer this range constraints.
        bits: usize,
    },

    /// A range `from`-`to`, both inclusive.
    Signed{
        /// Lower bound (inclusive).
        from: i64,
        /// Upper bound (inclusive).
        to: i64,
        /// Bitsize of the integer this range constraints.
        bits: usize,
    },

    /// A range containing all `bits` long integer.
    Full{
        /// Bitsize of the integer this range constraints.
        bits: usize,
    },
}

impl Constraint {
    /// Creates a new full ranges, `bits` large.
    pub fn new(bits: usize) -> Result<Constraint> {
        if bits == 0 { return Err("Constraints can't have bits size zero".into()); }
        Ok(Constraint::Full{ bits: bits })
    }

    /// Returns true if the range includes no values.
    pub fn is_empty(&self) -> bool {
        if let &Constraint::Empty{ .. } = self { true } else { false }
    }

    /// Sets the lower bound to the unsigned value `bound` (inclusive).
    pub fn clamp_lower_bound_unsigned(&mut self, bound: u64) {
        let bits = self.bits();
        self.make_unsigned();

        match self {
            &mut Constraint::Empty{ .. } => { /* do nothing */ }
            &mut Constraint::Unsigned{ ref mut from, ref mut to,.. } if *to >= bound =>{
                *from = cmp::max(bound,*from);
                *to = cmp::max(bound,*to);
            }
            a@&mut Constraint::Unsigned{ .. } => {
                *a = Constraint::Empty{ bits: bits };
            }
            &mut Constraint::Signed{ .. } => { unreachable!() }
            a@&mut Constraint::Full{ .. } => {
                *a = Constraint::Unsigned{
                    from: bound,
                    to: u64::MAX >> (64 - cmp::min(64,bits)),
                    bits: bits,
                };
            }
        }
    }

    /// Sets the lower bound to the signed value `bound` (inclusive).
    pub fn clamp_lower_bound_signed(&mut self, bound: u64) {
        let bits = self.bits();
        let missing = u64::MAX << (cmp::min(64,bits) % 64);
        let sign_bit = 1 << (cmp::min(64,bits) - 1);
        let bound = if sign_bit & bound != 0 { (bound | missing) as i64 } else { bound as i64 };

        self.make_signed();

        match self {
            &mut Constraint::Empty{ .. } => { /* do nothing */ }
            &mut Constraint::Unsigned{ .. } => { unreachable!() }
            &mut Constraint::Signed{ ref mut from, ref mut to,.. } if *to < bound =>{
                *from = cmp::max(bound,*from);
                *to = cmp::max(bound,*to);
            }
            a@&mut Constraint::Signed{ .. } => {
                *a = Constraint::Empty{ bits: bits };
            }
            a@&mut Constraint::Full{ .. } => {
                *a = Constraint::Signed{
                    from: bound,
                    to: i64::MAX >> (64 - cmp::min(64,bits)),
                    bits: bits,
                };
            }
        }
    }

    /// Sets the upper bound to the unsigned value `bound` (inclusive).
    pub fn clamp_upper_bound_unsigned(&mut self, bound: u64) {
        let bits = self.bits();
        self.make_unsigned();

        match self {
            &mut Constraint::Empty{ .. } => { /* do nothing */ }
            &mut Constraint::Unsigned{ ref mut from, ref mut to,.. } if *from < bound =>{
                *from = cmp::min(bound,*from);
                *to = cmp::min(bound,*to);
            }
            a@&mut Constraint::Unsigned{ .. } => {
                *a = Constraint::Empty{ bits: bits };
            }
            &mut Constraint::Signed{ .. } => { unreachable!() }
            a@&mut Constraint::Full{ .. } => {
                *a = Constraint::Unsigned{
                    from: 0,
                    to: bound,
                    bits: bits,
                };
            }
        }
    }

    /// Sets the upper bound to the signed value `bound` (inclusive).
    pub fn clamp_upper_bound_signed(&mut self, bound: u64) {
        let bits = self.bits();
        let missing = u64::MAX << (cmp::min(64,bits) % 64);
        let sign_bit = 1 << (cmp::min(64,bits) - 1);
        let bound = if sign_bit & bound != 0 { (bound | missing) as i64 } else { bound as i64 };

        self.make_signed();

        match self {
            &mut Constraint::Empty{ .. } => { /* do nothing */ }
            &mut Constraint::Unsigned{ .. } => { unreachable!() }
            &mut Constraint::Signed{ ref mut from, ref mut to,.. } if *from < bound =>{
                *from = cmp::min(bound,*from);
                *to = cmp::min(bound,*to);
            }
            a@&mut Constraint::Signed{ .. } => {
                *a = Constraint::Empty{ bits: bits };
            }
            a@&mut Constraint::Full{ .. } => {
                *a = Constraint::Signed{
                    from: i64::MIN >> (64 - cmp::min(64,bits)),
                    to: bound,
                    bits: bits,
                };
            }
        }
    }

    /// Extends the range to include all values in `other`.
    pub fn union_with(&mut self, mut other: Self) {
        match (self.clone(),other.clone()) {
            (Constraint::Empty{ .. },_) => { /* do nothing */ }
            (_,Constraint::Empty{ bits }) => {
                *self = Constraint::Empty{ bits: bits };
            }
            (_,Constraint::Full{ .. }) => { /* do nothing */ }
            (Constraint::Full{ .. },a) => {
                *self = a;
            }
            (Constraint::Signed{ from: from_s,.. },Constraint::Unsigned{ .. }) => {
                if from_s >= 0 {
                    self.make_unsigned();
                } else {
                    other.make_signed();
                }

                self.union_with(other);
            }
            (Constraint::Unsigned{ .. },Constraint::Signed{ from: from_s,.. }) => {
                if from_s >= 0 {
                    other.make_unsigned();
                } else {
                    self.make_signed();
                }

                self.union_with(other);
            }
            (Constraint::Unsigned{ from: from_a, to: to_a,.. },
             Constraint::Unsigned{ from: from_b, to: to_b, bits: bits_b }) => {
                let max = u64::MAX >> 64.saturating_sub(bits_b);
                let from = cmp::min(from_a as u64,from_b);
                let to = cmp::max(to_a as u64,to_b);

                *self = if from <= to && to <= max {
                    Constraint::Unsigned{
                        from: from,
                        to: to,
                        bits: bits_b,
                    }
                } else {
                    Constraint::Empty{ bits: bits_b }
                };
            }
            (Constraint::Signed{ from: from_a, to: to_a,.. },
             Constraint::Signed{ from: from_b, to: to_b, bits: bits_b }) => {
                let max = (u64::MAX >> 64.saturating_sub(bits_b) + 1) as i64;
                let from = cmp::min(from_a,from_b);
                let to = cmp::max(to_a,to_b);

                *self = if from <= to && to <= max {
                    Constraint::Signed{
                        from: from,
                        to: to,
                        bits: bits_b,
                    }
                } else {
                    Constraint::Empty{ bits: bits_b }
                };
            }
        }
    }

    /// Limits this range to the values in `other`.
    pub fn intersect_with(&mut self, mut other: Self) {
        match (self.clone(),other.clone()) {
            (Constraint::Empty{ .. },_) => { /* do nothing */ }
            (_,Constraint::Empty{ bits }) => {
                *self = Constraint::Empty{ bits: bits };
            }
            (_,Constraint::Full{ .. }) => { /* do nothing */ }
            (Constraint::Full{ .. },a) => {
                *self = a;
            }
            (Constraint::Signed{ from: from_s,.. },Constraint::Unsigned{ .. }) |
            (Constraint::Unsigned{ .. },Constraint::Signed{ from: from_s,.. }) => {
                if from_s >= 0 {
                    other.make_unsigned();
                } else {
                    other.make_signed();
                }

                self.union_with(other);
            }
            (Constraint::Unsigned{ from: from_a, to: to_a,.. },
             Constraint::Unsigned{ from: from_b, to: to_b, bits: bits_b }) => {
                let max = u64::MAX >> 64.saturating_sub(bits_b);
                let from = cmp::max(from_a as u64,from_b);
                let to = cmp::min(to_a as u64,to_b);

                *self = if from <= to && to <= max {
                    Constraint::Unsigned{
                        from: from,
                        to: to,
                        bits: bits_b,
                    }
                } else {
                    Constraint::Empty{ bits: bits_b }
                };
            }
            (Constraint::Signed{ from: from_a, to: to_a,.. },
             Constraint::Signed{ from: from_b, to: to_b, bits: bits_b }) => {
                let max = (u64::MAX >> 64.saturating_sub(bits_b) + 1) as i64;
                let from = cmp::max(from_a,from_b);
                let to = cmp::min(to_a,to_b);

                *self = if from <= to && to <= max as i64 {
                    Constraint::Signed{
                        from: from,
                        to: to,
                        bits: bits_b,
                    }
                } else {
                    Constraint::Empty{ bits: bits_b }
                };
            }
        }
    }

    /// Limits the range to `bound`.
    pub fn include(&mut self, bound: u64) {
        match self.clone() {
            Constraint::Empty{ .. } => { /* skip */ },
            Constraint::Signed{ .. } =>  { /* skip */ }
            Constraint::Unsigned{ bits, from, to } if from <= bound && to >= bound => {
                *self = Constraint::Unsigned{
                    from: bound,
                    to: bound,
                    bits: bits
                };
            }
            Constraint::Unsigned{ bits,.. } => {
                *self = Constraint::Empty{ bits: bits }
            }
            Constraint::Full{ bits } => {
                *self = Constraint::Unsigned{
                    from: bound,
                    to: bound,
                    bits: bits
                };
            }
        }
    }

    /// Move the range to exclude `bound`.
    pub fn exclude(&mut self, bound: u64) {
        match self.clone() {
            Constraint::Empty{ .. } => { /* skip */ },
            Constraint::Signed{ .. } =>  { /* skip */ }
            Constraint::Full{ .. } =>  { /* skip */ }
            Constraint::Unsigned{ bits, from, to } if from == bound && to == bound => {
                *self = Constraint::Empty{ bits: bits }
            }
            Constraint::Unsigned{ bits, from, to } if from == bound => {
                *self = Constraint::Unsigned{
                    from: bound + 1,
                    to: to,
                    bits: bits
                };
            }
            Constraint::Unsigned{ bits, from, to } if to == bound => {
                *self = Constraint::Unsigned{
                    from: from,
                    to: bound - 1,
                    bits: bits
                };
            }
            Constraint::Unsigned{ .. } =>  { /* skip */ }
        }
    }

    /// Size of the range bounds in bits.
    pub fn bits(&self) -> usize {
        match self {
            &Constraint::Empty{ bits } => bits,
            &Constraint::Signed{ bits,.. } => bits,
            &Constraint::Unsigned{ bits,.. } => bits,
            &Constraint::Full{ bits } => bits,
        }
    }

    /// Apply this range to `v`, returning either `v` itself of `Undefined` if it is outside.
    pub fn limit(&self, v: Value) -> Value {
        match v {
            Value::Undefined => Value::Undefined,
            Value::Constant(Constant{ value, bits }) => {
                let svalue = Self::signed(value,bits);
                match self {
                    &Constraint::Empty{ .. } => Value::Undefined,
                    &Constraint::Unsigned{ ref from, ref to,.. } if *from <= value && *to >= value => {
                        Value::val(value,bits).unwrap()
                    }
                    &Constraint::Unsigned{ .. } => Value::Undefined,
                    &Constraint::Signed{ ref from, ref to,.. } if *from <= svalue && *to >= svalue => {
                        Value::val(value,bits).unwrap()
                    }
                    &Constraint::Signed{ .. } => Value::Undefined,
                    &Constraint::Full{ .. } => Value::val(value,bits).unwrap()
                }
            }
            val@Value::Variable(_) => val,
        }
    }

    /// Helper to convert `value` to a `bits` large signed integer.
    pub fn signed(value: u64, bits: usize) -> i64 {
        let missing = if bits < 64 { u64::MAX << bits } else { 0 };
        let sign_bit = if bits < 64 { 1 << (cmp::min(64,bits) - 1) } else { 0x8000000000000000 };
        if sign_bit & value != 0 { (value | missing) as i64 } else { value as i64 }
    }

    /// Inverts the ranges to include all values it excluded before.
    pub fn invert(&mut self) {
        match self.clone() {
            Constraint::Empty{ bits } => {
                *self = Constraint::Full{ bits: bits };
            }
            Constraint::Full{ bits } => {
                *self = Constraint::Empty{ bits: bits };
            }
            Constraint::Signed{ to, bits,.. } if to >= Self::signed_max(bits) => {
                *self = Constraint::Signed{
                    from: i64::MIN >> (64 - cmp::min(64,bits)),
                    to: to.saturating_sub(1),
                    bits: bits,
                };
            }
            Constraint::Signed{ from, to, bits } if from <= i64::MIN >> (64 - cmp::min(64,bits)) => {
                *self = Constraint::Signed{
                    from: to.saturating_add(1),
                    to: i64::MAX >> (64 - cmp::min(64,bits)),
                    bits: bits,
                };
            }
            Constraint::Unsigned{ from, to, bits } if to >= Self::unsigned_max(bits) => {
                *self = Constraint::Unsigned{
                    from: 0,
                    to: from.saturating_sub(1),
                    bits: bits,
                };
            }
            Constraint::Unsigned{ from, to, bits } if from == 0 => {
                *self = Constraint::Unsigned{
                    from: to.saturating_add(1),
                    to: u64::MAX >> (64 - cmp::min(64,bits)),
                    bits: bits,
                };
            }
            _ => {
                *self = Constraint::Full{ bits: self.bits() };
            }
        }
    }

    /// Returns the inverted range. See `invert`.
    pub fn inverted(&self) -> Self {
        let mut r = self.clone();
        r.invert();
        r
    }

    fn make_unsigned(&mut self) {
        if let Constraint::Signed{ from, to, bits } = self.clone() {
            let from = cmp::max(0,from) as u64;
            let to = cmp::max(0,to) as u64;

            *self = if from < to {
                Constraint::Unsigned{
                    from: from,
                    to: to,
                    bits: bits,
                }
            } else {
                Constraint::Empty{ bits: bits }
            };
        }
    }

    fn make_signed(&mut self) {
        if let Constraint::Unsigned{ from, to, bits } = self.clone() {
            let from = cmp::min(u64::MAX >> 1,from) as i64;
            let to = cmp::min(u64::MAX >> 1,to) as i64;

            *self = if from < to {
                Constraint::Signed{
                    from: from,
                    to: to,
                    bits: bits,
                }
            } else {
                Constraint::Empty{ bits: bits }
            };
        }
    }

    fn unsigned_max(b: usize) -> u64 {
        u64::MAX >> (64 - cmp::min(64,b))
    }

    fn signed_max(b: usize) -> i64 {
        i64::MAX >> (64 - cmp::min(64,b))
    }
}

impl Arbitrary for Constraint {
    fn arbitrary<G: Gen>(g: &mut G) -> Self {
        let bits = g.gen_range(0,65);

        match g.gen_range(0, 4) {
            0 => Constraint::Empty{ bits: bits },
            1 => {
                let a = g.gen_range(0,u64::MAX >> cmp::min(63,bits));
                let b = g.gen_range(0,u64::MAX >> cmp::min(63,bits));

                Constraint::Unsigned{
                    from: cmp::min(a,b),
                    to: cmp::max(a,b),
                    bits: bits
                }
            },
            2 => {
                let a = g.gen_range(i64::MIN >> cmp::min(63,bits),i64::MAX >> cmp::min(63,bits));
                let b = g.gen_range(i64::MIN >> cmp::min(63,bits),i64::MAX >> cmp::min(63,bits));

                Constraint::Signed{
                    from: cmp::min(a,b),
                    to: cmp::max(a,b),
                    bits: bits,
                }
            }
            3 => Constraint::Full{ bits: bits },
            _ => { unreachable!() }
        }
    }
}

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

    #[test]
    fn constraint_include() {
        let mut c1 = Constraint::new(16).unwrap();

        c1.include(42);
        assert_eq!(c1, Constraint::Unsigned{ from: 42, to: 42, bits: 16 });

        c1.include(41);
        assert_eq!(c1, Constraint::Empty{ bits: 16 });
    }

    #[test]
    fn constraint_lower_bound_unsigned() {
        let mut c1 = Constraint::new(16).unwrap();

        c1.clamp_lower_bound_unsigned(23);
        assert_eq!(c1, Constraint::Unsigned{ from: 23, to: 0xFFFF, bits: 16 });

        c1.clamp_lower_bound_unsigned(42);
        assert_eq!(c1, Constraint::Unsigned{ from: 42, to: 0xFFFF, bits: 16 });

        c1.clamp_lower_bound_unsigned(23);
        assert_eq!(c1, Constraint::Unsigned{ from: 42, to: 0xFFFF, bits: 16 });
    }

    #[test]
    fn constraint_lower_bound_signed() {
        let mut c1 = Constraint::new(16).unwrap();

        c1.clamp_lower_bound_signed(23);
        assert_eq!(c1, Constraint::Signed{ from: 23, to: 0x7FFF, bits: 16 });

        c1.clamp_lower_bound_unsigned(42);
        assert_eq!(c1, Constraint::Unsigned{ from: 42, to: 0x7FFF, bits: 16 });

        c1.clamp_lower_bound_unsigned(23);
        assert_eq!(c1, Constraint::Unsigned{ from: 42, to: 0x7FFF, bits: 16 });
    }

    #[test]
    fn constraint_lower_bound_mixed() {
        let mut c1 = Constraint::new(16).unwrap();

        // -23
        c1.clamp_lower_bound_signed(0xffe9);
        assert_eq!(c1, Constraint::Signed{ from: -23, to: 0x7FFF, bits: 16 });

        c1.clamp_lower_bound_unsigned(23);
        assert_eq!(c1, Constraint::Unsigned{ from: 23, to: 0x7FFF, bits: 16 });
    }
    #[test]
    fn constraint_upper_bound_unsigned() {
        let mut c1 = Constraint::new(16).unwrap();

        c1.clamp_upper_bound_unsigned(0x1337);
        assert_eq!(c1, Constraint::Unsigned{ from: 0, to: 0x1337, bits: 16 });

        c1.clamp_lower_bound_unsigned(0x100);
        assert_eq!(c1, Constraint::Unsigned{ from: 0x100, to: 0x1337, bits: 16 });

        c1.clamp_lower_bound_unsigned(0x1337);
        assert_eq!(c1, Constraint::Unsigned{ from: 0x1337, to: 0x1337, bits: 16 });
    }

    #[test]
    fn constraint_upper_bound_signed() {
        let mut c1 = Constraint::new(16).unwrap();

        c1.clamp_upper_bound_signed(42);
        assert_eq!(c1, Constraint::Signed{ from: -0x8000, to: 42, bits: 16 });

        c1.clamp_upper_bound_signed(23);
        assert_eq!(c1, Constraint::Signed{ from: -0x8000, to: 23, bits: 16 });

        c1.clamp_upper_bound_signed(42);
        assert_eq!(c1, Constraint::Signed{ from: -0x8000, to: 23, bits: 16 });
    }

    #[test]
    fn constraint_upper_bound_mixed() {
        let mut c1 = Constraint::new(16).unwrap();

        // -23
        c1.clamp_upper_bound_signed(0xffe9);
        assert_eq!(c1, Constraint::Signed{ from: -0x8000, to: -23, bits: 16 });

        c1.clamp_upper_bound_unsigned(23);
        assert_eq!(c1, Constraint::Empty{ bits: 16 });
    }

    // XXX: limit

    #[test]
    fn max() {
        assert_eq!(Constraint::unsigned_max(1), 1);
        assert_eq!(Constraint::unsigned_max(8), 0xff);
        assert_eq!(Constraint::unsigned_max(16), 0xffff);
        assert_eq!(Constraint::unsigned_max(32), 0xffffffff);
        assert_eq!(Constraint::unsigned_max(64), 0xffffffffffffffff);
        assert_eq!(Constraint::unsigned_max(128), 0xffffffffffffffff);

        assert_eq!(Constraint::signed_max(1), 0);
        assert_eq!(Constraint::signed_max(8), 0x7f);
        assert_eq!(Constraint::signed_max(16), 0x7fff);
        assert_eq!(Constraint::signed_max(32), 0x7fffffff);
        assert_eq!(Constraint::signed_max(64), 0x7fffffffffffffff);
        assert_eq!(Constraint::signed_max(128), 0x7fffffffffffffff);
    }

    #[test]
    fn constraint_invert() {
        assert_eq!(Constraint::Empty{ bits: 16 }.inverted(), Constraint::Full{ bits: 16 });
        assert_eq!(Constraint::Full{ bits: 16 }.inverted(), Constraint::Empty{ bits: 16 });
        assert_eq!(Constraint::Unsigned{ from: 0, to: 100, bits: 16 }.inverted(), Constraint::Unsigned{ from: 101, to: 0xffff, bits: 16 });
        assert_eq!(Constraint::Unsigned{ from: 101, to: 0xffff, bits: 16 }.inverted(), Constraint::Unsigned{ from: 0, to: 100, bits: 16 });
        assert_eq!(Constraint::Unsigned{ from: 0, to: 0, bits: 16 }.inverted(), Constraint::Unsigned{ from: 1, to: 0xffff, bits: 16 });
        assert_eq!(Constraint::Unsigned{ from: 1, to: 0xffff, bits: 16 }.inverted(), Constraint::Unsigned{ from: 0, to: 0, bits: 16 });
        assert_eq!(Constraint::Unsigned{ from: 1, to: 2, bits: 16 }.inverted(), Constraint::Full{ bits: 16 });
    }
}
