/*
 * Copyright (c) Meta Platforms, Inc. and affiliates.
 *
 * This source code is licensed under the MIT license found in the
 * LICENSE file in the root directory of this source tree.
 */

mod abstract_partition_test {
    use im::HashSet;
    use sparta::datatype::AbstractDomain;
    use sparta::datatype::AbstractPartition;
    use sparta::datatype::HashMapAbstractPartition;
    use sparta::datatype::HashSetAbstractDomain;
    use sparta::datatype::PatriciaTreeMapAbstractPartition;

    type Domain = HashSetAbstractDomain<String>;

    fn build_domain<const N: usize>(arr: [&str; N]) -> Domain {
        arr.iter().map(|str| str.to_string()).collect()
    }

    #[test]
    fn test_hmap_lattice_operations() {
        type Partition = HashMapAbstractPartition<String, Domain>;

        assert!(Partition::top().leq(&Partition::top()));
        assert!(!Partition::top().leq(&Partition::bottom()));
        assert!(Partition::bottom().leq(&Partition::top()));
        assert!(Partition::bottom().leq(&Partition::bottom()));

        assert!(Partition::bottom() == Partition::bottom());
        assert!(Partition::top() == Partition::top());
        assert!(Partition::bottom() != Partition::top());

        let mut p1 = Partition::bottom();
        p1.set("v1".to_owned(), build_domain(["a", "b"]));
        p1.set("v2".to_owned(), build_domain(["c"]));
        p1.set("v3".to_owned(), build_domain(["d", "e", "f"]));
        p1.set("v4".to_owned(), build_domain(["a", "f"]));
        assert_eq!(p1.len(), 4);

        let mut p2 = Partition::bottom();
        p2.set("v0".to_owned(), build_domain(["c", "f"]));
        p2.set("v2".to_owned(), build_domain(["c", "d"]));
        p2.set("v3".to_owned(), build_domain(["d", "e", "g", "h"]));
        assert_eq!(p2.len(), 3);

        assert!(Partition::bottom().leq(&p1));
        assert!(!p1.leq(&Partition::bottom()));
        assert!(!Partition::top().leq(&p1));
        assert!(p1.leq(&Partition::top()));
        assert!(!p1.leq(&p2));
        assert!(!p2.leq(&p1));

        assert!(p1 == p1);
        assert!(p1 != p2);

        let join = p1.clone().join(p2.clone());
        assert!(p1.leq(&join));
        assert!(p2.leq(&join));
        assert_eq!(join.len(), 5);
        assert!(join.get(&("v0".to_string())) == p2.get(&("v0".to_string())));
        assert!(join.get(&("v1".to_string())) == p1.get(&("v1".to_string())));

        assert!(
            join.get(&("v2".to_string())).set()
                == &HashSet::from(vec!["c".to_string(), "d".to_string()])
        );
        assert!(
            join.get(&("v3".to_string())).set()
                == &HashSet::from(vec![
                    "d".to_string(),
                    "e".to_string(),
                    "f".to_string(),
                    "g".to_string(),
                    "h".to_string()
                ])
        );

        assert!(join.get(&("v4".to_string())) == p1.get(&("v4".to_string())));

        assert!(p1.clone().join(Partition::top()).is_top());
        assert!(p1.clone().join(Partition::bottom()) == (p1));

        let meet = p1.clone().meet(p2.clone());
        assert!(meet.leq(&p1));
        assert!(meet.leq(&p2));
        assert_eq!(meet.len(), 2);

        assert!(meet.get(&("v2".to_string())).set() == &HashSet::from(vec!["c".to_string()]));
        assert!(
            meet.get(&("v3".to_string())).set()
                == &HashSet::from(vec!["d".to_string(), "e".to_string()])
        );

        assert!(p1.clone().meet(Partition::bottom()).is_bottom());
        assert!(p1.clone().meet(Partition::top()) == p1);
    }

    #[test]
    fn test_ptmap_lattice_operations() {
        type Partition = PatriciaTreeMapAbstractPartition<u32, Domain>;

        assert!(Partition::top().leq(&Partition::top()));
        assert!(!Partition::top().leq(&Partition::bottom()));
        assert!(Partition::bottom().leq(&Partition::top()));
        assert!(Partition::bottom().leq(&Partition::bottom()));

        assert!(Partition::bottom() == Partition::bottom());
        assert!(Partition::top() == Partition::top());
        assert!(Partition::bottom() != Partition::top());

        let mut p1 = Partition::bottom();
        p1.set(1, build_domain(["a", "b"]));
        p1.set(2, build_domain(["c"]));
        p1.set(3, build_domain(["d", "e", "f"]));
        p1.set(4, build_domain(["a", "f"]));
        assert_eq!(p1.len(), 4);

        let mut p2 = Partition::bottom();
        p2.set(0, build_domain(["c", "f"]));
        p2.set(2, build_domain(["c", "d"]));
        p2.set(3, build_domain(["d", "e", "g", "h"]));
        assert_eq!(p2.len(), 3);

        assert!(Partition::bottom().leq(&p1));
        assert!(!p1.leq(&Partition::bottom()));
        assert!(!Partition::top().leq(&p1));
        assert!(p1.leq(&Partition::top()));
        assert!(!p1.leq(&p2));
        assert!(!p2.leq(&p1));

        assert!(p1 == p1);
        assert!(p1 != p2);

        let join = p1.clone().join(p2.clone());
        assert!(p1.leq(&join));
        assert!(p2.leq(&join));
        assert_eq!(join.len(), 5);
        assert!(join.get(&0) == p2.get(&0));
        assert!(join.get(&1) == p1.get(&1));

        assert!(join.get(&2).set() == &HashSet::from(vec!["c".to_string(), "d".to_string()]));
        assert!(
            join.get(&3).set()
                == &HashSet::from(vec![
                    "d".to_string(),
                    "e".to_string(),
                    "f".to_string(),
                    "g".to_string(),
                    "h".to_string()
                ])
        );

        assert!(join.get(&4) == p1.get(&4));

        assert!(p1.clone().join(Partition::top()).is_top());
        assert!(p1.clone().join(Partition::bottom()) == (p1));

        let meet = p1.clone().meet(p2.clone());
        assert!(meet.leq(&p1));
        assert!(meet.leq(&p2));
        assert_eq!(meet.len(), 2);

        assert!(meet.get(&2).set() == &HashSet::from(vec!["c".to_string()]));
        assert!(meet.get(&3).set() == &HashSet::from(vec!["d".to_string(), "e".to_string()]));

        assert!(p1.clone().meet(Partition::bottom()).is_bottom());
        assert!(p1.clone().meet(Partition::top()) == p1);
    }
}


/*
 * Copyright (c) Meta Platforms, Inc. and affiliates.
 *
 * This source code is licensed under the MIT license found in the
 * LICENSE file in the root directory of this source tree.
 */

//! Tests borrowed from the MonotonicFixpointIteratorTest.cpp in the C++ version.

mod liveness {
    use std::borrow::Cow;
    use std::collections::HashMap;

    use im::HashSet;
    use smallvec::SmallVec;
    use sparta::datatype::HashSetAbstractDomain;
    use sparta::fixpoint_iter::FixpointIteratorTransformer;
    use sparta::fixpoint_iter::MonotonicFixpointIterator;
    use sparta::graph::Graph;
    use sparta::graph::ReverseGraph;
    use sparta::graph::ReversedRefGraph;
    use sparta::graph::DEFAULT_GRAPH_SUCCS_NUM;

    pub type NodeId = u32;
    pub type EdgeId = u32;

    pub struct Edge(NodeId, NodeId);

    pub type Symbol = &'static str;

    pub struct Statement {
        uses: Vec<Symbol>,
        defs: Vec<Symbol>,
    }

    impl Statement {
        pub fn new(uses: Vec<Symbol>, defs: Vec<Symbol>) -> Self {
            Self { uses, defs }
        }
    }

    #[derive(Default)]
    pub struct Program {
        statements: HashMap<NodeId, Statement>,
        edge_interner: Vec<Edge>,
        edges: HashMap<NodeId, HashSet<EdgeId>>,
        pred_edges: HashMap<NodeId, HashSet<EdgeId>>,
        stmt_index: NodeId,
        exit: NodeId,
    }

    impl Program {
        pub fn add_stmt(&mut self, stmt: Statement) -> NodeId {
            self.statements.insert(self.stmt_index, stmt);
            let cur_index = self.stmt_index;
            self.stmt_index += 1;
            cur_index
        }

        pub fn stmt_at(&self, n: NodeId) -> &Statement {
            self.statements
                .get(&n)
                .unwrap_or_else(|| panic!("Statement {} does not exist. ", n))
        }

        pub fn add_edge(&mut self, src: NodeId, dst: NodeId) {
            self.edge_interner.push(Edge(src, dst));
            let edge_id = self.edge_interner.len() as u32 - 1;
            self.edges.entry(src).or_default().insert(edge_id);
            self.pred_edges.entry(dst).or_default().insert(edge_id);
        }

        pub fn set_exit(&mut self, exit: NodeId) {
            self.exit = exit;
        }
    }

    impl Graph for Program {
        type NodeId = NodeId;
        type EdgeId = EdgeId;

        fn entry(&self) -> Self::NodeId {
            0
        }

        fn exit(&self) -> Self::NodeId {
            self.exit
        }

        fn predecessors(
            &self,
            n: Self::NodeId,
        ) -> SmallVec<[Self::EdgeId; DEFAULT_GRAPH_SUCCS_NUM]> {
            self.pred_edges
                .get(&n)
                .map(|v| v.iter().copied().collect())
                .unwrap_or_else(SmallVec::new)
        }

        fn successors(&self, n: Self::NodeId) -> SmallVec<[Self::EdgeId; DEFAULT_GRAPH_SUCCS_NUM]> {
            self.edges
                .get(&n)
                .map(|v| v.iter().copied().collect())
                .unwrap_or_else(SmallVec::new)
        }

        fn source(&self, e: Self::EdgeId) -> Self::NodeId {
            self.edge_interner[e as usize].0
        }

        fn target(&self, e: Self::EdgeId) -> Self::NodeId {
            self.edge_interner[e as usize].1
        }

        fn size(&self) -> usize {
            self.statements.len()
        }
    }

    type LivenessDomain = HashSetAbstractDomain<Symbol>;

    pub struct LivenessTransformer<'a> {
        graph: &'a Program,
    }

    impl<'a> LivenessTransformer<'a> {
        pub fn new(graph: &'a Program) -> Self {
            Self { graph }
        }
    }

    impl<'a> FixpointIteratorTransformer<ReversedRefGraph<'a, Program>, LivenessDomain>
        for LivenessTransformer<'a>
    {
        fn analyze_node(&mut self, n: NodeId, current_state: &mut LivenessDomain) {
            let stmt = self.graph.stmt_at(n);
            current_state.remove_elements(&stmt.defs);
            current_state.add_elements(stmt.uses.clone());
        }

        fn analyze_edge(
            &mut self,
            _: EdgeId,
            exit_state_at_src: &LivenessDomain,
        ) -> LivenessDomain {
            exit_state_at_src.clone()
        }
    }

    trait LivenessAnalysis {
        fn get_live_in_vars_at(&self, n: NodeId) -> Cow<'_, LivenessDomain>;
        fn get_live_out_vars_at(&self, n: NodeId) -> Cow<'_, LivenessDomain>;
    }

    type LivenessFixpointIterator<'g> = MonotonicFixpointIterator<
        'g,
        ReversedRefGraph<'g, Program>,
        LivenessDomain,
        LivenessTransformer<'g>,
    >;

    impl<'g> LivenessAnalysis for LivenessFixpointIterator<'g> {
        fn get_live_in_vars_at(&self, n: NodeId) -> Cow<'_, LivenessDomain> {
            self.get_exit_state_at(n)
        }

        fn get_live_out_vars_at(&self, n: NodeId) -> Cow<'_, LivenessDomain> {
            self.get_entry_state_at(n)
        }
    }

    /**
     *                       live in          live out
     *  1: a = 0;             {c}              {a, c}
     *  2: b = a + 1;         {a, c}           {b, c}
     *  3: c = c + b;         {b, c}           {b, c}
     *  4: a = b * 2;         {b, c}           {a, c}
     *  5: if (a < 9) {       {a, c}           {a, c}
     *       goto 2;
     *     } else {
     *  6:   return c;        {c}              {}
     *     }
     */
    fn build_program1() -> Program {
        let mut program = Program::default();
        let l1 = program.add_stmt(Statement::new(vec![], vec!["a"]));
        let l2 = program.add_stmt(Statement::new(vec!["a"], vec!["b"]));
        let l3 = program.add_stmt(Statement::new(vec!["c", "b"], vec!["c"]));
        let l4 = program.add_stmt(Statement::new(vec!["b"], vec!["a"]));
        let l5 = program.add_stmt(Statement::new(vec!["a"], vec![]));
        let l6 = program.add_stmt(Statement::new(vec!["c"], vec![]));
        program.add_edge(l1, l2);
        program.add_edge(l2, l3);
        program.add_edge(l3, l4);
        program.add_edge(l4, l5);
        program.add_edge(l5, l6);
        program.add_edge(l5, l2);
        program.set_exit(l6);
        program
    }

    /**
     *                       live in          live out
     *  1: x = a + b;        {a, b}           {x, a, b}
     *  2: y = a * b;        {x, a, b}        {x, y, a, b}
     *  3: if (y > a) {      {x, y, a, b}     {x, y, a, b}
     *  4:   return x;       {x}              {}
     *     }
     *  5: a = a + 1;        {y, a, b}        {y, a, b}
     *  6: x = a + b;        {y, a, b}        {x, y, a, b}
     *     if (...) {
     *       goto 7;
     *     }
     *     goto 3;
     *  7: x = y + a;
     *
     */
    fn build_program2() -> Program {
        let mut program = Program::default();
        let l1 = program.add_stmt(Statement::new(vec!["a", "b"], vec!["x"]));
        let l2 = program.add_stmt(Statement::new(vec!["a", "b"], vec!["y"]));
        let l3 = program.add_stmt(Statement::new(vec!["y", "a"], vec![]));
        let l4 = program.add_stmt(Statement::new(vec!["x"], vec![]));
        let l5 = program.add_stmt(Statement::new(vec!["a"], vec!["a"]));
        let l6 = program.add_stmt(Statement::new(vec!["a", "b"], vec!["x"]));
        let l7 = program.add_stmt(Statement::new(vec!["y", "a"], vec!["x"]));
        program.add_edge(l1, l2);
        program.add_edge(l2, l3);
        program.add_edge(l3, l4);
        program.add_edge(l3, l5);
        program.add_edge(l5, l6);
        program.add_edge(l6, l3);
        program.add_edge(l6, l7);
        program.set_exit(l4);
        program
    }

    /*
     *                          live in           live out
     *  1: a, b -> x, y         {a, b, z}         {a, b, x, y, z}
     *  2: x, y -> z            {x, y, a, b}      {a, b, y, z}
     *  3: a -> c               {a, b, y, z}      {c, b, y, z}
     *  4: b -> d               {c, b, y, z}      {c, d, y, z}
     *  5: c, d -> a, b         {c, d, y, z}      {a, b, y, z}
     *  6: a, b -> x            {a, b, y, z}      {a, b, x, y, z}
     *  7: return z             {z}               {}
     *  8: a, b -> c, d         {a, b, y, z}      {c, b, y, z}
     *
     *  1->2, 2->3, 3->4, 4->5, 5->6, 6->7, 6->2, 5->3, 1->8, 8->4
     */
    fn build_program3() -> Program {
        let mut program = Program::default();
        let l1 = program.add_stmt(Statement::new(vec!["a", "b"], vec!["x", "y"]));
        let l2 = program.add_stmt(Statement::new(vec!["x", "y"], vec!["z"]));
        let l3 = program.add_stmt(Statement::new(vec!["a"], vec!["c"]));
        let l4 = program.add_stmt(Statement::new(vec!["b"], vec!["d"]));
        let l5 = program.add_stmt(Statement::new(vec!["c", "d"], vec!["a", "b"]));
        let l6 = program.add_stmt(Statement::new(vec!["a", "b"], vec!["x"]));
        let l7 = program.add_stmt(Statement::new(vec!["z"], vec![]));
        let l8 = program.add_stmt(Statement::new(vec!["a", "b"], vec!["c", "d"]));
        program.add_edge(l1, l2);
        program.add_edge(l2, l3);
        program.add_edge(l3, l4);
        program.add_edge(l4, l5);
        program.add_edge(l5, l6);
        program.add_edge(l6, l7);
        program.add_edge(l6, l2);
        program.add_edge(l5, l3);
        program.add_edge(l1, l8);
        program.add_edge(l8, l4);
        program.set_exit(l7);
        program
    }

    macro_rules! assert_analysis_values {
        ( $fp:ident, $index:literal, $method:ident, [$($vars:literal),*] ) => {
            let mut elements: Vec<Symbol> = $fp.$method($index).elements().copied().collect();
            elements.sort();
            let expected: Vec<Symbol> = vec![$($vars),*];
            assert_eq!(elements, expected);
        };
    }

    macro_rules! assert_analysis {
        ( $fp:ident, $index:literal, [$($live_in_vars:literal),*], [$($live_out_vars:literal),*] ) => {
            assert!(matches!(
                $fp.get_exit_state_at($index).as_ref(),
                &LivenessDomain::Value(_)
            ));
            assert!(matches!(
                $fp.get_entry_state_at($index).as_ref(),
                &LivenessDomain::Value(_)
            ));
            assert_analysis_values!($fp, $index, get_live_in_vars_at, [$($live_in_vars),*]);
            assert_analysis_values!($fp, $index, get_live_out_vars_at, [$($live_out_vars),*]);
        };
    }

    #[test]
    fn test_fixpoint_iter_liveness_program1() {
        let program = build_program1();
        let reversed_program = program.rev();
        let transformer = LivenessTransformer::new(&program);
        let mut fp =
            MonotonicFixpointIterator::new(&reversed_program, 4, transformer, &reversed_program);
        fp.run(LivenessDomain::value_from_set(HashSet::new()));

        // 0: a = 0;
        assert_analysis!(fp, 0, ["c"], ["a", "c"]);

        // 1: b = a + 1;
        assert_analysis!(fp, 1, ["a", "c"], ["b", "c"]);

        // 2: c = c + b;
        assert_analysis!(fp, 2, ["b", "c"], ["b", "c"]);

        // 3: a = b * 2;
        assert_analysis!(fp, 3, ["b", "c"], ["a", "c"]);

        // 4: if (a < 9)
        assert_analysis!(fp, 4, ["a", "c"], ["a", "c"]);

        // 5: return c;
        assert_analysis!(fp, 5, ["c"], []);
    }

    #[test]
    fn test_fixpoint_iter_liveness_program2() {
        let program = build_program2();
        let reversed_program = program.rev();
        let transformer = LivenessTransformer::new(&program);
        let mut fp =
            MonotonicFixpointIterator::new(&reversed_program, 4, transformer, &reversed_program);
        fp.run(LivenessDomain::value_from_set(HashSet::new()));

        // 0: x = a + b;
        assert_analysis!(fp, 0, ["a", "b"], ["a", "b", "x"]);

        // 1: y = a * b;
        assert_analysis!(fp, 1, ["a", "b", "x"], ["a", "b", "x", "y"]);

        // 2: if (y > a);
        assert_analysis!(fp, 2, ["a", "b", "x", "y"], ["a", "b", "x", "y"]);

        // 3: return x;
        assert_analysis!(fp, 3, ["x"], []);

        // 4: a = a + 1;
        assert_analysis!(fp, 4, ["a", "b", "y"], ["a", "b", "y"]);

        // 5: x = a + b;
        assert_analysis!(fp, 5, ["a", "b", "y"], ["a", "b", "x", "y"]);

        // 7: x = y + a;
        assert!(matches!(
            fp.get_exit_state_at(6).as_ref(),
            &LivenessDomain::Bottom
        ));
        assert!(matches!(
            fp.get_entry_state_at(6).as_ref(),
            &LivenessDomain::Bottom
        ));
    }

    #[test]
    fn test_fixpoint_iter_liveness_program3() {
        let program = build_program3();
        let reversed_program = program.rev();
        let transformer = LivenessTransformer::new(&program);
        let mut fp =
            MonotonicFixpointIterator::new(&reversed_program, 4, transformer, &reversed_program);
        fp.run(LivenessDomain::value_from_set(HashSet::new()));

        // 0: a, b -> x, y
        assert_analysis!(fp, 0, ["a", "b", "z"], ["a", "b", "x", "y", "z"]);

        // 1: x, y -> z
        assert_analysis!(fp, 1, ["a", "b", "x", "y"], ["a", "b", "y", "z"]);

        // 2: a -> c
        assert_analysis!(fp, 2, ["a", "b", "y", "z"], ["b", "c", "y", "z"]);

        // 3: b -> d
        assert_analysis!(fp, 3, ["b", "c", "y", "z"], ["c", "d", "y", "z"]);

        // 4: c, d -> a, b
        assert_analysis!(fp, 4, ["c", "d", "y", "z"], ["a", "b", "y", "z"]);

        // 5: a, b -> x
        assert_analysis!(fp, 5, ["a", "b", "y", "z"], ["a", "b", "x", "y", "z"]);

        // 7: return z
        assert_analysis!(fp, 6, ["z"], []);

        // 8: a, b -> c, d
        assert_analysis!(fp, 7, ["a", "b", "y", "z"], ["b", "c", "y", "z"]);
    }
}

mod numerical {
    use smallvec::SmallVec;
    use sparta::datatype::AbstractDomain;
    use sparta::datatype::AbstractEnvironment;
    use sparta::datatype::PatriciaTreeMapAbstractEnvironment;
    use sparta::datatype::PatriciaTreeSet;
    use sparta::datatype::SetAbstractDomainOps;
    use sparta::fixpoint_iter::FixpointIteratorTransformer;
    use sparta::fixpoint_iter::MonotonicFixpointIterator;
    use sparta::graph::Graph;
    use sparta::graph::DEFAULT_GRAPH_SUCCS_NUM;

    type NodeId = usize;
    type EdgeId = usize;

    struct Edge(NodeId, NodeId);

    // For simplicity, we directly use u32 to refer variable.
    type Variable = u32;

    #[derive(Clone, Copy)]
    enum Statement {
        Assignment {
            var: Variable,
            value: u32,
        },
        Addition {
            result: Variable,
            lhs: Variable,
            rhs: u32,
        },
    }

    #[derive(Default)]
    struct BasicBlock {
        statements: Vec<Statement>,
        successors: Vec<EdgeId>,
        predecessors: Vec<EdgeId>,
    }

    impl BasicBlock {
        fn add(&mut self, stmt: Statement) {
            self.statements.push(stmt);
        }

        fn statements(&self) -> &[Statement] {
            self.statements.as_slice()
        }

        fn successors(&self) -> &[EdgeId] {
            self.successors.as_slice()
        }

        fn predecessors(&self) -> &[EdgeId] {
            self.predecessors.as_slice()
        }
    }

    #[derive(Default)]
    struct Program {
        basic_blocks: Vec<BasicBlock>,
        edges: Vec<Edge>,
        entry: NodeId,
        exit: NodeId,
    }

    impl Program {
        fn create_block(&mut self) -> NodeId {
            let id = self.basic_blocks.len();
            let bb = BasicBlock::default();
            self.basic_blocks.push(bb);
            id
        }

        fn get_block_mut(&mut self, n: NodeId) -> &mut BasicBlock {
            &mut self.basic_blocks[n]
        }

        fn add_edge(&mut self, source: NodeId, target: NodeId) {
            let edge_id = self.edges.len();
            self.edges.push(Edge(source, target));
            self.basic_blocks[source].successors.push(edge_id);
            self.basic_blocks[target].predecessors.push(edge_id);
        }

        fn set_entry(&mut self, entry: NodeId) {
            self.entry = entry;
        }

        fn set_exit(&mut self, exit: NodeId) {
            self.exit = exit;
        }
    }

    impl Graph for Program {
        type NodeId = NodeId;
        type EdgeId = EdgeId;

        fn entry(&self) -> Self::NodeId {
            self.entry
        }

        fn exit(&self) -> Self::NodeId {
            self.exit
        }

        fn predecessors(
            &self,
            n: Self::NodeId,
        ) -> SmallVec<[Self::EdgeId; DEFAULT_GRAPH_SUCCS_NUM]> {
            self.basic_blocks[n]
                .predecessors()
                .iter()
                .copied()
                .collect()
        }

        fn successors(&self, n: Self::NodeId) -> SmallVec<[Self::EdgeId; DEFAULT_GRAPH_SUCCS_NUM]> {
            self.basic_blocks[n].successors().iter().copied().collect()
        }

        fn source(&self, e: Self::EdgeId) -> Self::NodeId {
            self.edges[e].0
        }

        fn target(&self, e: Self::EdgeId) -> Self::NodeId {
            self.edges[e].1
        }

        fn size(&self) -> usize {
            self.basic_blocks.len()
        }
    }

    #[derive(Debug, Clone, PartialEq, Eq)]
    enum IntegerSetAbstractDomain {
        Top,
        Value(PatriciaTreeSet<u32>),
    }

    impl IntegerSetAbstractDomain {
        fn insert(&mut self, value: u32) {
            if let Self::Value(set) = self {
                set.insert(value);
            }
        }

        fn add(lhs: &Self, rhs: &Self) -> Self {
            match (lhs, rhs) {
                (Self::Top, _) | (_, Self::Top) => Self::Top,
                (Self::Value(lset), Self::Value(rset)) => {
                    if lset.is_empty() || rset.is_empty() {
                        Self::bottom()
                    } else {
                        let mut res = Self::bottom();
                        for x in lset {
                            for y in rset {
                                res.insert(x + y);
                            }
                        }
                        res
                    }
                }
            }
        }
    }

    impl AbstractDomain for IntegerSetAbstractDomain {
        fn bottom() -> Self {
            IntegerSetAbstractDomain::Value(PatriciaTreeSet::new())
        }

        fn top() -> Self {
            IntegerSetAbstractDomain::Top
        }

        fn is_bottom(&self) -> bool {
            if let Self::Value(set) = self {
                return set.is_empty();
            }
            false
        }

        fn is_top(&self) -> bool {
            matches!(self, Self::Top)
        }

        fn leq(&self, rhs: &Self) -> bool {
            match (self, rhs) {
                (_, Self::Top) => true,
                (Self::Top, _) => false,
                (Self::Value(lset), Self::Value(rset)) => {
                    if lset.is_empty() {
                        true
                    } else if rset.is_empty() {
                        false
                    } else {
                        lset.is_subset(rset)
                    }
                }
            }
        }

        fn join_with(&mut self, rhs: Self) {
            match (self, rhs) {
                (Self::Top, _) => {}
                (lhs @ _, Self::Top) => {
                    *lhs = Self::Top;
                }
                (Self::Value(lset), Self::Value(rset)) => {
                    if rset.is_empty() {
                        return;
                    }
                    if lset.is_empty() {
                        *lset = rset;
                    } else {
                        lset.union_with(rset);
                    }
                }
            }
        }

        fn meet_with(&mut self, _rhs: Self) {}

        fn widen_with(&mut self, rhs: Self) {
            match (&self, rhs) {
                (Self::Top, _) => {}
                (_, Self::Top) => {
                    *self = Self::Top;
                }
                (Self::Value(lset), Self::Value(rset)) => {
                    if rset.is_subset(&lset) {
                        return;
                    }
                    *self = Self::Top;
                }
            }
        }

        fn narrow_with(&mut self, _rhs: Self) {}
    }

    type IntegerSetAbstractEnvironment =
        PatriciaTreeMapAbstractEnvironment<u32, IntegerSetAbstractDomain>;

    pub struct IntegerSetTransformer<'a> {
        prog: &'a Program,
    }

    impl<'a> IntegerSetTransformer<'a> {
        fn new(prog: &'a Program) -> Self {
            Self { prog }
        }

        fn analyze_statement(&self, stmt: &Statement, env: &mut IntegerSetAbstractEnvironment) {
            match *stmt {
                Statement::Assignment { var, value } => {
                    env.set(var, IntegerSetAbstractDomain::Value([value].into()));
                }
                Statement::Addition { result, lhs, rhs } => {
                    env.set(
                        result,
                        IntegerSetAbstractDomain::add(
                            &env.get(&lhs).into_owned(),
                            &IntegerSetAbstractDomain::Value([rhs].into()),
                        ),
                    );
                }
            }
        }
    }

    impl<'a> FixpointIteratorTransformer<Program, IntegerSetAbstractEnvironment>
        for IntegerSetTransformer<'a>
    {
        fn analyze_node(&mut self, n: NodeId, env: &mut IntegerSetAbstractEnvironment) {
            let bb = &self.prog.basic_blocks[n];
            for stmt in bb.statements() {
                self.analyze_statement(stmt, env);
            }
        }

        fn analyze_edge(
            &mut self,
            _: EdgeId,
            env: &IntegerSetAbstractEnvironment,
        ) -> IntegerSetAbstractEnvironment {
            env.clone()
        }
    }

    /**
     * bb1: x = 1;
     *      if (...) {
     * bb2:   y = x + 1;
     *      } else {
     * bb3:   y = x + 2;
     *      }
     * bb4: return
     */
    fn build_program1() -> Program {
        let mut program = Program::default();
        let bb1 = program.create_block();
        let bb2 = program.create_block();
        let bb3 = program.create_block();
        let bb4 = program.create_block();

        let x = 0u32; // "x"
        let y = 1u32; // "y"
        program
            .get_block_mut(bb1)
            .add(Statement::Assignment { var: x, value: 1 });
        program.add_edge(bb1, bb2);
        program.add_edge(bb1, bb3);

        program.get_block_mut(bb2).add(Statement::Addition {
            result: y,
            lhs: x,
            rhs: 1,
        });
        program.add_edge(bb2, bb4);

        program.get_block_mut(bb3).add(Statement::Addition {
            result: y,
            lhs: x,
            rhs: 2,
        });
        program.add_edge(bb3, bb4);

        program.set_entry(bb1);
        program.set_exit(bb4);

        program
    }

    /**
     * bb1: x = 1;
     *      while (...) {
     * bb2:   x = x + 1;
     *      }
     * bb3: return
     */
    fn build_program2() -> Program {
        let mut program = Program::default();
        let bb1 = program.create_block();
        let bb2 = program.create_block();
        let bb3 = program.create_block();

        let x = 0u32; // "x"
        program
            .get_block_mut(bb1)
            .add(Statement::Assignment { var: x, value: 1 });
        program.add_edge(bb1, bb2);

        program.get_block_mut(bb2).add(Statement::Addition {
            result: x,
            lhs: x,
            rhs: 1,
        });
        program.add_edge(bb2, bb2);
        program.add_edge(bb2, bb3);

        program.set_entry(bb1);
        program.set_exit(bb3);

        program
    }

    #[test]
    fn test_fixpoint_iter_integerset_program1() {
        let prog = build_program1();

        let transformer = IntegerSetTransformer::new(&prog);
        let mut fp = MonotonicFixpointIterator::new(&prog, 4, transformer, &prog);
        fp.run(IntegerSetAbstractEnvironment::top());

        let x = 0u32;
        let y = 1u32;

        let bb1 = 0;
        assert_eq!(
            fp.get_entry_state_at(bb1).into_owned(),
            IntegerSetAbstractEnvironment::top()
        );
        assert_eq!(
            fp.get_exit_state_at(bb1).get(&x).into_owned(),
            IntegerSetAbstractDomain::Value([1].into())
        );
        assert_eq!(
            fp.get_exit_state_at(bb1).get(&y).into_owned(),
            IntegerSetAbstractDomain::top()
        );

        let bb2 = 1;
        assert_eq!(fp.get_entry_state_at(bb2), fp.get_exit_state_at(bb1));
        assert_eq!(
            fp.get_exit_state_at(bb2).get(&x).into_owned(),
            IntegerSetAbstractDomain::Value([1].into())
        );
        assert_eq!(
            fp.get_exit_state_at(bb2).get(&y).into_owned(),
            IntegerSetAbstractDomain::Value([2].into())
        );

        let bb3 = 2;
        assert_eq!(fp.get_entry_state_at(bb3), fp.get_exit_state_at(bb1));
        assert_eq!(
            fp.get_exit_state_at(bb3).get(&x).into_owned(),
            IntegerSetAbstractDomain::Value([1].into())
        );
        assert_eq!(
            fp.get_exit_state_at(bb3).get(&y).into_owned(),
            IntegerSetAbstractDomain::Value([3].into())
        );

        let bb4 = 3;
        assert_eq!(fp.get_entry_state_at(bb4), fp.get_exit_state_at(bb4));
        assert_eq!(
            fp.get_exit_state_at(bb4).get(&x).into_owned(),
            IntegerSetAbstractDomain::Value([1].into())
        );
        assert_eq!(
            fp.get_exit_state_at(bb4).get(&y).into_owned(),
            IntegerSetAbstractDomain::Value([2, 3].into())
        );
    }

    #[test]
    fn test_fixpoint_iter_integerset_program2() {
        let prog = build_program2();

        let transformer = IntegerSetTransformer::new(&prog);
        let mut fp = MonotonicFixpointIterator::new(&prog, 4, transformer, &prog);
        fp.run(IntegerSetAbstractEnvironment::top());

        let x = 0u32;

        let bb1 = 0;
        assert_eq!(
            fp.get_entry_state_at(bb1).into_owned(),
            IntegerSetAbstractEnvironment::top()
        );
        assert_eq!(
            fp.get_exit_state_at(bb1).get(&x).into_owned(),
            IntegerSetAbstractDomain::Value([1].into())
        );

        let bb2 = 1;
        assert_eq!(
            fp.get_entry_state_at(bb2).get(&x).into_owned(),
            IntegerSetAbstractDomain::top()
        );
        assert_eq!(
            fp.get_exit_state_at(bb2).get(&x).into_owned(),
            IntegerSetAbstractDomain::top()
        );

        let bb3 = 2;
        assert_eq!(
            fp.get_entry_state_at(bb3).get(&x).into_owned(),
            IntegerSetAbstractDomain::top()
        );
        assert_eq!(
            fp.get_exit_state_at(bb3).get(&x).into_owned(),
            IntegerSetAbstractDomain::top()
        );
    }
}


/*
 * Copyright (c) Meta Platforms, Inc. and affiliates.
 *
 * This source code is licensed under the MIT license found in the
 * LICENSE file in the root directory of this source tree.
 */

mod common;

use common::graph::SimpleGraph;
use sparta::graph::ReverseGraph;
use sparta::graph::SuccessorNodes;

fn create_graph() -> SimpleGraph {
    let mut g = SimpleGraph::new(1, 5);
    g.add_edge(1, 2);
    g.add_edge(1, 3);
    g.add_edge(3, 4);
    g.add_edge(2, 5);
    g.add_edge(4, 5);
    g
}

#[test]
fn test_graph_successors() {
    let g = create_graph();

    let mut succ_nodes = g.get_succ_nodes(1);
    succ_nodes.sort();
    assert_eq!(succ_nodes.into_vec(), vec![2, 3]);

    let mut succ_nodes = g.get_succ_nodes(4);
    succ_nodes.sort();
    assert_eq!(succ_nodes.into_vec(), vec![5]);
}

#[test]
fn test_reversed_graph() {
    let graph = create_graph();
    let reversed_g = graph.into_rev();

    let mut succ_nodes = reversed_g.get_succ_nodes(5);
    succ_nodes.sort();
    assert_eq!(succ_nodes.into_vec(), vec![2, 4]);

    let mut succ_nodes = reversed_g.get_succ_nodes(1);
    succ_nodes.sort();
    assert_eq!(succ_nodes.into_vec(), vec![]);
}


/*
 * Copyright (c) Meta Platforms, Inc. and affiliates.
 *
 * This source code is licensed under the MIT license found in the
 * LICENSE file in the root directory of this source tree.
 */

mod common;

use std::collections::HashMap;

use common::graph::NodeId;
use common::graph::SimpleGraph;
use sparta::graph::Graph;
use sparta::wpo::WeakPartialOrdering;
use sparta::wpo::WpoIdx;

// (NodeId in Graph, is plain node?, is head node?, is exit node?,
//  number of successors, number of predecessors, number of outer predecessors).
type NodeProperty = (NodeId, bool, bool, bool, usize, u32, usize);

fn check_wpo(wpo: WeakPartialOrdering<NodeId>, expected_results: &[NodeProperty], node_seq: &str) {
    let mut count = HashMap::<WpoIdx, u32>::new();
    let mut worklist = Vec::new();
    worklist.push(wpo.get_entry());
    let mut output = String::new();
    let mut first = true;
    let mut i = 0;
    while let Some(v) = worklist.pop() {
        for &w in wpo.get_successors(v).iter() {
            let c = count.entry(w).or_default();
            *c += 1;
            if *c == wpo.get_num_preds(w) {
                worklist.push(w);
            }
        }

        let expected = expected_results[i];
        i += 1;
        assert_eq!(expected.0, wpo.get_node(v));
        assert_eq!(expected.1, wpo.is_plain(v));
        assert_eq!(expected.2, wpo.is_head(v));
        assert_eq!(expected.3, wpo.is_exit(v));
        assert_eq!(expected.4, wpo.get_successors(v).len());
        assert_eq!(expected.5, wpo.get_num_preds(v));
        if wpo.is_head(v) {
            assert_eq!(wpo.get_node(wpo.get_exit_of_head(v)), wpo.get_node(v));
            if !first {
                output.push(' ');
            }
            output.push('(');
            output.push_str(&wpo.get_node(v).to_string());
        } else if wpo.is_exit(v) {
            assert_eq!(expected.6, wpo.get_num_outer_preds(v).len());
            assert_eq!(wpo.get_node(wpo.get_head_of_exit(v)), wpo.get_node(v));
            output.push(')');
        } else {
            if !first {
                output.push(' ');
            }
            output.push_str(&wpo.get_node(v).to_string());
        }

        first = false;
    }
    assert_eq!(output, node_seq);
}

/*
 * This graph and the corresponding weak partial ordering are described
 * on page 4 of Bourdoncle's paper:
 *   F. Bourdoncle. Efficient chaotic iteration strategies with widenings.
 *   In Formal Methods in Programming and Their Applications, pp 128-141.
 * The graph is given as follows:
 *
 *                 +-----------------------+
 *                 |           +-----+     |
 *                 |           |     |     |
 *                 V           V     |     |
 *     1 --> 2 --> 3 --> 4 --> 5 --> 6 --> 7 --> 8
 *           |           |                 ^     ^
 *           |           |                 |     |
 *           |           +-----------------+     |
 *           +-----------------------------------+
 *
 * Bourdoncle's algorithm computes the following weak partial ordering:
 *
 *     1 2 (3 4 (5 6) 7) 8
 */
#[test]
fn test_wpo_example_from_wto_paper() {
    let mut g = SimpleGraph::new(1, 8);
    g.add_edge(1, 2);
    g.add_edge(2, 3);
    g.add_edge(3, 4);
    g.add_edge(4, 5);
    g.add_edge(5, 6);
    g.add_edge(6, 7);
    g.add_edge(7, 8);
    g.add_edge(2, 8);
    g.add_edge(4, 7);
    g.add_edge(6, 5);
    g.add_edge(7, 3);
    // "1 2 (3 4 (5 6) 7) 8"

    println!("Size: {}", g.size());
    let wpo = WeakPartialOrdering::new(1, g.size(), &g);

    assert!(!wpo.is_from_outside(5, 6));
    assert!(!wpo.is_from_outside(3, 7));
    assert!(wpo.is_from_outside(3, 2));
    assert!(!wpo.is_from_outside(3, 4));

    assert_eq!(wpo.size(), 10);

    let expected_results = [
        (1, true, false, false, 1, 0, 0),
        (2, true, false, false, 1, 1, 0),
        (3, false, true, false, 1, 1, 0),
        (4, true, false, false, 1, 1, 0),
        (5, false, true, false, 1, 1, 0),
        (6, true, false, false, 1, 1, 0),
        (5, false, false, true, 1, 1, 1),
        (7, true, false, false, 1, 1, 0),
        (3, false, false, true, 1, 1, 1),
        (8, true, false, false, 0, 1, 0),
    ];

    check_wpo(wpo, &expected_results, "1 2 (3 4 (5 6) 7) 8");
}

#[test]
fn test_wpo_example_from_wpo_paper() {
    let mut g = SimpleGraph::new(1, 10);
    g.add_edge(1, 2);
    g.add_edge(2, 3);
    g.add_edge(3, 4);
    g.add_edge(4, 3);
    g.add_edge(3, 5);
    g.add_edge(5, 2);
    g.add_edge(2, 6);
    g.add_edge(6, 5);
    g.add_edge(6, 7);
    g.add_edge(7, 8);
    g.add_edge(8, 6);
    g.add_edge(6, 9);
    g.add_edge(9, 8);
    g.add_edge(2, 10);

    let wpo = WeakPartialOrdering::new(1, g.size(), &g);

    assert_eq!(13, wpo.size());

    let expected_results = [
        (1, true, false, false, 1, 0, 0),
        (2, false, true, false, 2, 1, 0),
        (3, false, true, false, 1, 1, 0),
        (4, true, false, false, 1, 1, 0),
        (3, false, false, true, 1, 1, 1),
        (6, false, true, false, 2, 1, 0),
        (7, true, false, false, 1, 1, 0),
        (9, true, false, false, 1, 1, 0),
        (8, true, false, false, 1, 2, 0),
        (6, false, false, true, 1, 1, 1),
        (5, true, false, false, 1, 2, 0),
        (2, false, false, true, 1, 1, 1),
        (10, true, false, false, 0, 1, 0),
    ];

    check_wpo(wpo, &expected_results, "1 (2 (3 4) (6 7 9 8) 5) 10");
}

#[test]
fn test_wpo_single_node() {
    let g = SimpleGraph::new(1, 1);
    let wpo = WeakPartialOrdering::new(1, 1, &g);

    assert_eq!(1, wpo.size());

    let expected_results = [(1, true, false, false, 0, 0, 0)];

    check_wpo(wpo, &expected_results, "1");
}

//             +--+
//             v  |
// +---+     +------+
// | 1 | --> |  2   |
// +---+     +------+
#[test]
fn test_wpo_single_scc_at_end() {
    let mut g = SimpleGraph::new(1, 2);
    g.add_edge(1, 2);
    g.add_edge(2, 2);

    let wpo = WeakPartialOrdering::new(1, g.size(), &g);

    assert_eq!(3, wpo.size());

    let expected_results = [
        (1, true, false, false, 1, 0, 0),
        (2, false, true, false, 1, 1, 0),
        (2, false, false, true, 0, 1, 1),
    ];

    check_wpo(wpo, &expected_results, "1 (2)");
}

//             +--+
//             v  |
// +---+     +------+     +---+
// | 1 | <-- |  2   | --> | 3 |
// +---+     +------+     +---+
//   |         ^
//   +---------+
#[test]
fn test_wpo_single_scc_at_end_of_scc() {
    let mut g = SimpleGraph::new(1, 3);
    g.add_edge(1, 2);
    g.add_edge(2, 1);
    g.add_edge(2, 2);
    g.add_edge(2, 3);

    let wpo = WeakPartialOrdering::new(1, g.size(), &g);

    assert_eq!(5, wpo.size());
    assert!(!wpo.is_from_outside(2, 2));
    assert!(!wpo.is_from_outside(1, 2));
    assert!(wpo.is_from_outside(2, 1));

    let expected_results = [
        (1, false, true, false, 1, 0, 0),
        (2, false, true, false, 1, 1, 0),
        (2, false, false, true, 1, 1, 1),
        (1, false, false, true, 1, 1, 0),
        (3, true, false, false, 0, 1, 0),
    ];

    check_wpo(wpo, &expected_results, "(1 (2)) 3");
}

//             +---------+
//             v         |
// +---+     +---+     +---+
// | 1 | --> | 2 | --> | 3 |
// +---+     +---+     +---+
#[test]
fn test_wpo_scc_at_end() {
    let mut g = SimpleGraph::new(1, 3);
    g.add_edge(1, 2);
    g.add_edge(2, 3);
    g.add_edge(3, 2);

    let wpo = WeakPartialOrdering::new(1, g.size(), &g);

    assert!(!wpo.is_from_outside(2, 3));
    assert!(wpo.is_from_outside(2, 1));
    assert_eq!(4, wpo.size());

    let expected_results = [
        (1, true, false, false, 1, 0, 0),
        (2, false, true, false, 1, 1, 0),
        (3, true, false, false, 1, 1, 0),
        (2, false, false, true, 0, 1, 1),
    ];

    check_wpo(wpo, &expected_results, "1 (2 3)");
}

//   +-------------------+
//   |                   v
// +---+     +---+     +---+     +---+
// | 2 | <-- | 1 | <-- | 3 | --> | 4 |
// +---+     +---+     +---+     +---+
//   ^                   |
//   +-------------------+
#[test]
fn test_wpo_scc_at_end_of_scc() {
    let mut g = SimpleGraph::new(1, 4);
    g.add_edge(1, 2);
    g.add_edge(2, 3);
    g.add_edge(3, 2);
    g.add_edge(3, 1);
    g.add_edge(3, 4);

    let wpo = WeakPartialOrdering::new(1, g.size(), &g);

    assert!(!wpo.is_from_outside(1, 3));
    assert!(!wpo.is_from_outside(2, 3));
    assert!(wpo.is_from_outside(2, 1));
    assert_eq!(6, wpo.size());

    let expected_results = [
        (1, false, true, false, 1, 0, 0),
        (2, false, true, false, 1, 1, 0),
        (3, true, false, false, 1, 1, 0),
        (2, false, false, true, 1, 1, 1),
        (1, false, false, true, 1, 1, 0),
        (4, true, false, false, 0, 1, 0),
    ];

    check_wpo(wpo, &expected_results, "(1 (2 3)) 4");
}

#[test]
fn test_wpo_example_from_wpo_paper_irreducible() {
    let mut g = SimpleGraph::new(1, 6);
    g.add_edge(1, 2);
    g.add_edge(2, 3);
    g.add_edge(3, 2);
    g.add_edge(3, 4);
    g.add_edge(4, 3);
    g.add_edge(2, 5);
    g.add_edge(5, 4);
    g.add_edge(1, 6);
    g.add_edge(6, 4);

    let wpo = WeakPartialOrdering::new(1, g.size(), &g);

    assert!(!wpo.is_from_outside(3, 4));
    assert!(!wpo.is_from_outside(2, 3));
    assert!(wpo.is_from_outside(2, 6));
    assert!(wpo.is_from_outside(3, 6));
    assert!(wpo.is_from_outside(3, 5));
    assert!(!wpo.is_from_outside(2, 5));
    assert_eq!(8, wpo.size());

    let expected_results = [
        (1, true, false, false, 2, 0, 0),
        (2, false, true, false, 2, 1, 0),
        (3, false, true, false, 1, 1, 0),
        (5, true, false, false, 1, 1, 0),
        (6, true, false, false, 1, 1, 0),
        (4, true, false, false, 1, 3, 0),
        (3, false, false, true, 1, 1, 2),
        (2, false, false, true, 0, 1, 2),
    ];

    check_wpo(wpo, &expected_results, "1 (2 (3 5 6 4))");
}

#[test]
fn test_wpo_handle_outer_preds1() {
    let mut g = SimpleGraph::new(1, 93);
    g.add_edge(1, 12);
    g.add_edge(1, 16);
    g.add_edge(1, 18);
    g.add_edge(1, 26);
    g.add_edge(12, 45);
    g.add_edge(12, 75);
    g.add_edge(12, 46);
    g.add_edge(16, 74);
    g.add_edge(16, 75);
    g.add_edge(18, 92);
    g.add_edge(26, 93);
    g.add_edge(45, 46);
    g.add_edge(46, 47);
    g.add_edge(47, 73);
    g.add_edge(73, 74);
    g.add_edge(73, 75);
    g.add_edge(73, 73);
    g.add_edge(74, 46);
    g.add_edge(75, 45);
    g.add_edge(92, 93);
    g.add_edge(93, 45);
    g.add_edge(93, 46);

    let wpo = WeakPartialOrdering::new(1, g.size(), &g);

    assert_eq!(16, wpo.size());

    let expected_results = [
        (1, true, false, false, 4, 0, 0),
        (12, true, false, false, 1, 1, 0),
        (16, true, false, false, 2, 1, 0),
        (18, true, false, false, 1, 1, 0),
        (92, true, false, false, 1, 1, 0),
        (26, true, false, false, 1, 1, 0),
        (93, true, false, false, 2, 2, 0),
        (45, false, true, false, 1, 2, 0),
        (46, false, true, false, 1, 2, 0),
        (47, true, false, false, 1, 1, 0),
        (73, false, true, false, 1, 1, 0),
        (73, false, false, true, 1, 1, 1),
        (74, true, false, false, 1, 2, 0),
        (46, false, false, true, 1, 1, 2),
        (75, true, false, false, 1, 2, 0),
        (45, false, false, true, 0, 1, 4),
    ];

    check_wpo(
        wpo,
        &expected_results,
        "1 12 16 18 92 26 93 (45 (46 47 (73) 74) 75)",
    );
}

#[test]
fn test_wpo_handle_outer_preds2() {
    let mut g = SimpleGraph::new(1, 8);
    g.add_edge(1, 2);
    g.add_edge(2, 3);
    g.add_edge(3, 4);
    g.add_edge(4, 5);
    g.add_edge(5, 6);
    g.add_edge(6, 7);
    g.add_edge(6, 2);
    g.add_edge(5, 3);
    g.add_edge(1, 8);
    g.add_edge(8, 3);

    let wpo = WeakPartialOrdering::new(1, g.size(), &g);

    assert_eq!(10, wpo.size());

    let expected_results = [
        (1, true, false, false, 2, 0, 0),
        (2, false, true, false, 1, 1, 0),
        (8, true, false, false, 1, 1, 0),
        (3, false, true, false, 1, 2, 0),
        (4, true, false, false, 1, 1, 0),
        (5, true, false, false, 1, 1, 0),
        (3, false, false, true, 1, 1, 1),
        (6, true, false, false, 1, 1, 0),
        (2, false, false, true, 1, 1, 2),
        (7, true, false, false, 0, 1, 0),
    ];

    check_wpo(wpo, &expected_results, "1 (2 8 (3 4 5) 6) 7");
}

#[test]
fn test_wpo_handle_outer_preds3() {
    let mut g = SimpleGraph::new(1, 8);
    g.add_edge(1, 2);
    g.add_edge(2, 3);
    g.add_edge(3, 4);
    g.add_edge(4, 5);
    g.add_edge(5, 6);
    g.add_edge(6, 7);
    g.add_edge(6, 2);
    g.add_edge(5, 3);
    g.add_edge(1, 8);
    g.add_edge(8, 4);
    g.add_edge(7, 1);

    let wpo = WeakPartialOrdering::new(1, g.size(), &g);

    assert_eq!(11, wpo.size());

    let expected_results = [
        (1, false, true, false, 2, 0, 0),
        (2, false, true, false, 1, 1, 0),
        (3, false, true, false, 1, 1, 0),
        (8, true, false, false, 1, 1, 0),
        (4, true, false, false, 1, 2, 0),
        (5, true, false, false, 1, 1, 0),
        (3, false, false, true, 1, 1, 2),
        (6, true, false, false, 1, 1, 0),
        (2, false, false, true, 1, 1, 2),
        (7, true, false, false, 1, 1, 0),
        (1, false, false, true, 0, 1, 0),
    ];

    check_wpo(wpo, &expected_results, "(1 (2 (3 8 4 5) 6) 7)");
}

#[test]
fn test_wpo_handle_outer_preds4() {
    let mut g = SimpleGraph::new(0, 9);
    g.add_edge(0, 1);
    g.add_edge(1, 2);
    g.add_edge(2, 3);
    g.add_edge(3, 4);
    g.add_edge(4, 5);
    g.add_edge(5, 6);
    g.add_edge(6, 7);
    g.add_edge(6, 2);
    g.add_edge(5, 3);
    g.add_edge(1, 8);
    g.add_edge(8, 4);
    g.add_edge(7, 1);
    g.add_edge(7, 9);

    let wpo = WeakPartialOrdering::new(0, g.size(), &g);

    assert_eq!(13, wpo.size());

    let expected_results = [
        (0, true, false, false, 1, 0, 0),
        (1, false, true, false, 2, 1, 0),
        (2, false, true, false, 1, 1, 0),
        (3, false, true, false, 1, 1, 0),
        (8, true, false, false, 1, 1, 0),
        (4, true, false, false, 1, 2, 0),
        (5, true, false, false, 1, 1, 0),
        (3, false, false, true, 1, 1, 2),
        (6, true, false, false, 1, 1, 0),
        (2, false, false, true, 1, 1, 2),
        (7, true, false, false, 1, 1, 0),
        (1, false, false, true, 1, 1, 1),
        (9, true, false, false, 0, 1, 0),
    ];

    check_wpo(wpo, &expected_results, "0 (1 (2 (3 8 4 5) 6) 7) 9");
}

#[test]
fn test_wpo_handle_nested_loops_with_branch() {
    let mut g = SimpleGraph::new(1, 5);
    g.add_edge(1, 2);
    g.add_edge(1, 4);
    g.add_edge(2, 3);
    g.add_edge(3, 2);
    g.add_edge(3, 5);
    g.add_edge(4, 5);
    g.add_edge(5, 1);

    let wpo = WeakPartialOrdering::new(1, g.size(), &g);

    assert_eq!(7, wpo.size());

    let expected_results = [
        (1, false, true, false, 2, 0, 0),
        (2, false, true, false, 1, 1, 0),
        (3, true, false, false, 1, 1, 0),
        (2, false, false, true, 1, 1, 1),
        (4, true, false, false, 1, 1, 0),
        (5, true, false, false, 1, 2, 0),
        (1, false, false, true, 0, 1, 0),
    ];

    check_wpo(wpo, &expected_results, "(1 (2 3) 4 5)");
}


/*
 * Copyright (c) Meta Platforms, Inc. and affiliates.
 *
 * This source code is licensed under the MIT license found in the
 * LICENSE file in the root directory of this source tree.
 */

use std::collections::BTreeMap;
use std::collections::BTreeSet;
use std::collections::HashMap;
use std::collections::HashSet;
use std::iter::FromIterator;

use smallvec::SmallVec;
use sparta::graph::Graph;
use sparta::graph::DEFAULT_GRAPH_SUCCS_NUM;

pub type NodeId = u32;
pub type EdgeId = u32;

pub struct Edge(NodeId, NodeId);

// A naive graph for testing.
#[derive(Default)]
pub struct SimpleGraph {
    // We must use BTreeSet here to make the test consistent (we assumed
    // the order of successor edges in WPO tests.)
    edges: HashMap<NodeId, BTreeSet<EdgeId>>,
    pred_edges: BTreeMap<NodeId, BTreeSet<EdgeId>>,
    edge_interner: Vec<Edge>,
    enter: NodeId,
    exit: NodeId,
}

impl SimpleGraph {
    pub fn new(enter: NodeId, exit: NodeId) -> Self {
        Self {
            enter,
            exit,
            ..Default::default()
        }
    }

    pub fn add_edge(&mut self, source: NodeId, target: NodeId) {
        // Duplicate edges will be counted too.
        self.edge_interner.push(Edge(source, target));
        let edge_id = self.edge_interner.len() - 1;
        self.edges.entry(source).or_default().insert(edge_id as u32);
        self.pred_edges
            .entry(target)
            .or_default()
            .insert(edge_id as u32);
    }
}

impl Graph for SimpleGraph {
    type NodeId = NodeId;
    type EdgeId = EdgeId;

    fn entry(&self) -> Self::NodeId {
        self.enter
    }

    fn exit(&self) -> Self::NodeId {
        self.exit
    }

    // The edge is in the format of (source, target), so the predecessor
    // is still edge.0 instead of edge.1.
    fn predecessors(&self, n: Self::NodeId) -> SmallVec<[Self::EdgeId; DEFAULT_GRAPH_SUCCS_NUM]> {
        self.pred_edges
            .get(&n)
            .map(|v| v.iter().copied().collect())
            .unwrap_or_else(SmallVec::new)
    }

    fn successors(&self, n: Self::NodeId) -> SmallVec<[Self::EdgeId; DEFAULT_GRAPH_SUCCS_NUM]> {
        self.edges
            .get(&n)
            .map(|v| v.iter().copied().collect())
            .unwrap_or_else(SmallVec::new)
    }

    fn source(&self, e: Self::EdgeId) -> Self::NodeId {
        self.edge_interner[e as usize].0
    }

    fn target(&self, e: Self::EdgeId) -> Self::NodeId {
        self.edge_interner[e as usize].1
    }

    fn size(&self) -> usize {
        // Ideally the graph structure should maintain some stats of nodes and edges,
        // but for this testing code let's keep it simple. Note that we only provide
        // function `add_edge` to record edges, we don't keep information of nodes
        // separately, so this function will not give correct answer for graph with only
        // nodes but not edges.
        let mut nodes: HashSet<NodeId> = HashSet::from_iter(self.edges.keys().copied());
        nodes.extend(self.pred_edges.keys().copied());
        nodes.len()
    }
}


/*
 * Copyright (c) Meta Platforms, Inc. and affiliates.
 *
 * This source code is licensed under the MIT license found in the
 * LICENSE file in the root directory of this source tree.
 */

pub mod graph;


/*
 * Copyright (c) Meta Platforms, Inc. and affiliates.
 *
 * This source code is licensed under the MIT license found in the
 * LICENSE file in the root directory of this source tree.
 */

use std::hash::Hash;

use sparta::datatype::AbstractDomain;
use sparta::datatype::DisjointUnion;
use sparta::datatype::HashSetAbstractDomain;

#[derive(Clone, DisjointUnion, PartialEq, Eq)]
enum MyUnionedDomain {
    FirstCase(HashSetAbstractDomain<i32>),
    SecondCase(HashSetAbstractDomain<i64>),
}

#[test]
fn test_basic_union_top_bottom_leq() {
    assert!(MyUnionedDomain::top().is_top());
    assert!(MyUnionedDomain::SecondCase(AbstractDomain::top()).is_top());
    assert!(MyUnionedDomain::SecondCase(AbstractDomain::bottom()).is_bottom());

    let top: MyUnionedDomain = AbstractDomain::top();
    let bot: MyUnionedDomain = AbstractDomain::bottom();

    assert!(bot.leq(&top));

    let hsdom1: HashSetAbstractDomain<_> = [1, 2].into_iter().collect();
    let hsdom2: HashSetAbstractDomain<_> = [1, 2, 3].into_iter().collect();

    assert!(hsdom1.leq(&hsdom2));

    assert!(MyUnionedDomain::FirstCase(hsdom1).leq(&MyUnionedDomain::FirstCase(hsdom2)));
}

#[test]
fn test_diff_arms_no_leq() {
    let hsdom1: HashSetAbstractDomain<i32> = [1, 2].into_iter().collect();
    let hsdom2: HashSetAbstractDomain<i64> = [1, 2, 3].into_iter().collect();

    // watch out! Two different arms should not leq one another.
    let mudom1 = MyUnionedDomain::FirstCase(hsdom1);
    let mudom2 = MyUnionedDomain::SecondCase(hsdom2);

    assert!(!mudom1.leq(&mudom2));
    assert!(!mudom2.leq(&mudom1));
}

#[test]
fn test_join_same_arm() {
    let hsdom1: HashSetAbstractDomain<_> = [1, 2].into_iter().collect();
    let hsdom2: HashSetAbstractDomain<_> = [2, 3].into_iter().collect();

    let joined = hsdom1.clone().join(hsdom2.clone());
    let expected_joined: HashSetAbstractDomain<_> = [1, 2, 3].into_iter().collect();

    assert_eq!(joined, expected_joined);

    let mudom1 = MyUnionedDomain::FirstCase(hsdom1);
    let mudom2 = MyUnionedDomain::FirstCase(hsdom2);

    let joined_mudom = mudom1.join(mudom2);
    let inner_dom = match joined_mudom {
        MyUnionedDomain::FirstCase(inner) => inner,
        _ => panic!("Unexpected case"),
    };

    assert_eq!(inner_dom, expected_joined);
}

#[test]
fn test_meet_same_arm() {
    let hsdom1: HashSetAbstractDomain<_> = [1, 2].into_iter().collect();
    let hsdom2: HashSetAbstractDomain<_> = [2, 3].into_iter().collect();

    let met = hsdom1.clone().meet(hsdom2.clone());
    let expected_met: HashSetAbstractDomain<_> = [2].into_iter().collect();

    assert_eq!(met, expected_met);

    let mudom1 = MyUnionedDomain::FirstCase(hsdom1);
    let mudom2 = MyUnionedDomain::FirstCase(hsdom2);

    let met_mudom = mudom1.meet(mudom2);
    let inner_dom = match met_mudom {
        MyUnionedDomain::FirstCase(inner) => inner,
        _ => panic!("Unexpected case"),
    };

    assert_eq!(inner_dom, expected_met);
}

#[test]
fn test_join_diff_arm() {
    let hsdom1: HashSetAbstractDomain<_> = [1, 2].into_iter().collect();
    let hsdom2: HashSetAbstractDomain<_> = [2, 3].into_iter().collect();

    let mudom1 = MyUnionedDomain::FirstCase(hsdom1);
    let mudom2 = MyUnionedDomain::SecondCase(hsdom2);

    let joined_mudom = mudom1.join(mudom2);

    assert!(joined_mudom.is_top());
}

#[test]
fn test_meet_diff_arm() {
    let hsdom1: HashSetAbstractDomain<_> = [1, 2].into_iter().collect();
    let hsdom2: HashSetAbstractDomain<_> = [2, 3].into_iter().collect();

    let mudom1 = MyUnionedDomain::FirstCase(hsdom1);
    let mudom2 = MyUnionedDomain::SecondCase(hsdom2);

    let met_mudom = mudom1.meet(mudom2);

    assert!(met_mudom.is_bottom());
}

#[derive(Clone, DisjointUnion, PartialEq, Eq)]
enum TestGenericsDeriveTypechecks<S, T>
where
    S: Eq + Hash + Clone,
    T: Eq + Hash + Clone,
{
    FirstCase(HashSetAbstractDomain<S>),
    SecondCase(HashSetAbstractDomain<T>),
}

#[derive(Clone, DisjointUnion, PartialEq, Eq)]
enum TestGenericsDeriveForWholeDomainTypechecks<S: AbstractDomain, T>
where
    T: AbstractDomain,
{
    FirstCase(S),
    SecondCase(T),
}


/*
 * Copyright (c) Meta Platforms, Inc. and affiliates.
 *
 * This source code is licensed under the MIT license found in the
 * LICENSE file in the root directory of this source tree.
 */

#![crate_type = "proc-macro"]

use quote::quote;
use syn::parse_macro_input;
use syn::DeriveInput;

extern crate proc_macro;
use proc_macro::TokenStream;

#[proc_macro_derive(DisjointUnion)]
pub fn derive_disjoint_union(input: TokenStream) -> TokenStream {
    let code_item: DeriveInput = parse_macro_input!(input as DeriveInput);
    // Expect to be enum.
    let data_enum = match code_item.data {
        syn::Data::Enum(data_enum) => data_enum,
        _ => panic!("#[derive(DisjointUnion)] must be applied to enums only."),
    };

    let variant_idents = data_enum
        .variants
        .iter()
        .map(|variant| &variant.ident)
        .collect::<Vec<_>>();

    assert!(
        !variant_idents.is_empty(),
        "enum needs at least one variant."
    );

    let first_variant = variant_idents.first().unwrap();

    let enum_name = &code_item.ident;
    let (impl_generics, ty_generics, where_clause) = code_item.generics.split_for_impl();

    let expanded = quote! {
        impl #impl_generics AbstractDomain for #enum_name #ty_generics #where_clause {
            fn bottom() -> Self {
                #enum_name::#first_variant(AbstractDomain::bottom())
            }

            fn top() -> Self {
                #enum_name::#first_variant(AbstractDomain::top())
            }

            fn is_bottom(&self) -> bool {
                match self {
                    #( #enum_name::#variant_idents(dom) => dom.is_bottom(), )*
                }
            }

            fn is_top(&self) -> bool {
                match self {
                    #( #enum_name::#variant_idents(dom) => dom.is_top(), )*
                }
            }

            fn leq(&self, rhs: &Self) -> bool {
                if self.is_bottom() {
                    return true;
                }

                if rhs.is_bottom() {
                    return false;
                }

                if rhs.is_top() {
                    return true;
                }

                if self.is_top() {
                    return false;
                }

                return match (self, rhs) {
                    #( (#enum_name::#variant_idents(ref ldom), #enum_name::#variant_idents(rdom)) => ldom.leq(rdom), )*
                    _ => false,
                };
            }

            fn join_with(&mut self, rhs: Self) {
                match (self, rhs) {
                    #( (#enum_name::#variant_idents(ref mut ldom), #enum_name::#variant_idents(rdom)) => ldom.join_with(rdom), )*
                    (s, _) => *s = Self::top(),
                }
            }

            fn meet_with(&mut self, rhs: Self) {
                match (self, rhs) {
                    #( (#enum_name::#variant_idents(ref mut ldom), #enum_name::#variant_idents(rdom)) => ldom.meet_with(rdom), )*
                    (s, _) => *s = Self::bottom(),
                }
            }

            fn widen_with(&mut self, rhs: Self) {
                self.join_with(rhs)
            }

            fn narrow_with(&mut self, rhs: Self) {
                self.meet_with(rhs)
            }
        }
    };

    TokenStream::from(expanded)
}


/*
 * Copyright (c) Meta Platforms, Inc. and affiliates.
 *
 * This source code is licensed under the MIT license found in the
 * LICENSE file in the root directory of this source tree.
 */

use std::time::Instant;

use colored::*;
use log::debug;

use super::perf_logger::PerfLogEvent;
use super::perf_logger::PerfLogger;

pub struct ConsoleLogEvent;

pub struct ConsoleLogger;

impl PerfLogEvent for ConsoleLogEvent {
    type Timer = (&'static str, Instant);
    fn number(&self, name: &'static str, number: usize) {
        debug!("{}: {}", name, number);
    }
    fn string(&self, name: &'static str, value: String) {
        debug!("{}: {}", name, value);
    }
    fn start(&self, name: &'static str) -> Self::Timer {
        (name, Instant::now())
    }
    fn stop(&self, timer: Self::Timer) {
        let (name, time) = timer;
        print_time(name, time);
    }
    fn complete(self) {}
}

impl PerfLogger for ConsoleLogger {
    type PerfLogEvent = ConsoleLogEvent;
    fn create_event(&self, _name: &'static str) -> Self::PerfLogEvent {
        ConsoleLogEvent
    }
}

pub fn print_time(name: &str, time: Instant) {
    let elapsed_ms = time.elapsed().as_millis();
    let elapsed_str = format!("{:4}ms", elapsed_ms);
    let elapsed_color = if elapsed_ms < 10 {
        elapsed_str.dimmed()
    } else if elapsed_ms < 100 {
        elapsed_str.blue()
    } else if elapsed_ms < 1000 {
        elapsed_str.bold()
    } else {
        elapsed_str.red()
    };
    debug!("{} {}", elapsed_color, name.dimmed());
}


/*
 * Copyright (c) Meta Platforms, Inc. and affiliates.
 *
 * This source code is licensed under the MIT license found in the
 * LICENSE file in the root directory of this source tree.
 */

use std::collections::BTreeMap;
use std::error::Error;
use std::fmt;
use std::fmt::Write;

use lsp_types::DiagnosticSeverity;
use lsp_types::DiagnosticTag;
use serde_json::Value;

use crate::Location;
use crate::SourceLocationKey;

pub type Diagnostics = Vec<Diagnostic>;
pub type DiagnosticsResult<T> = Result<T, Diagnostics>;

#[derive(fmt::Debug)]
pub struct WithDiagnostics<T> {
    pub item: T,
    pub diagnostics: Vec<Diagnostic>,
}

impl<T> From<WithDiagnostics<T>> for Result<T, Diagnostics> {
    fn from(s: WithDiagnostics<T>) -> Result<T, Diagnostics> {
        if s.diagnostics.is_empty() {
            Ok(s.item)
        } else {
            Err(s.diagnostics)
        }
    }
}

pub fn diagnostics_result<T>(result: T, diagnostics: Diagnostics) -> DiagnosticsResult<T> {
    if diagnostics.is_empty() {
        Ok(result)
    } else {
        Err(diagnostics)
    }
}

/// A diagnostic message as a result of validating some code. This struct is
/// modeled after the LSP Diagnostic type:
/// https://microsoft.github.io/language-server-protocol/specification#diagnostic
///
/// Changes from LSP:
/// - `location` is different from LSP in that it's a file + span instead of
///   just a span.
/// - Unused fields are omitted.
#[derive(fmt::Debug)]
pub struct Diagnostic(Box<DiagnosticData>);

impl Diagnostic {
    fn with_severity<T: 'static + DiagnosticDisplay>(
        severity: DiagnosticSeverity,
        message: T,
        location: Location,
        tags: Vec<DiagnosticTag>,
    ) -> Self {
        Self(Box::new(DiagnosticData {
            message: Box::new(message),
            location,
            related_information: Vec::new(),
            tags,
            severity,
            data: Vec::new(),
            machine_readable: BTreeMap::new(),
        }))
    }

    /// Creates a new error Diagnostic.
    /// Additional locations can be added with the `.annotate()` function.
    pub fn error<T: 'static + DiagnosticDisplay>(message: T, location: Location) -> Self {
        Diagnostic::with_severity(DiagnosticSeverity::ERROR, message, location, Vec::new())
    }

    /// Creates a new error Diagnostic with additional data that
    /// can be used in IDE code actions
    pub fn error_with_data<T: 'static + DiagnosticDisplay + WithDiagnosticData>(
        message: T,
        location: Location,
    ) -> Self {
        let data = message.get_data();
        Self(Box::new(DiagnosticData {
            message: Box::new(message),
            location,
            tags: Vec::new(),
            severity: DiagnosticSeverity::ERROR,
            related_information: Vec::new(),
            data,
            machine_readable: BTreeMap::new(),
        }))
    }

    /// Creates a new Diagnostic with a severity of Warning
    /// Additional locations can be added with the `.annotate()` function.
    pub fn warning<T: 'static + DiagnosticDisplay>(
        message: T,
        location: Location,
        tags: Vec<DiagnosticTag>,
    ) -> Self {
        Diagnostic::with_severity(DiagnosticSeverity::WARNING, message, location, tags)
    }

    /// Creates a new Diagnostic with a severity of Information
    /// Additional locations can be added with the `.annotate()` function.
    pub fn info<T: 'static + DiagnosticDisplay>(
        message: T,
        location: Location,
        tags: Vec<DiagnosticTag>,
    ) -> Self {
        Diagnostic::with_severity(DiagnosticSeverity::INFORMATION, message, location, tags)
    }

    /// Creates a new Diagnostic with a severity of Hint
    /// Additional locations can be added with the `.annotate()` function.
    pub fn hint<T: 'static + DiagnosticDisplay>(
        message: T,
        location: Location,
        tags: Vec<DiagnosticTag>,
    ) -> Self {
        Diagnostic::with_severity(DiagnosticSeverity::HINT, message, location, tags)
    }

    /// Annotates this error with an additional location and associated message.
    pub fn annotate<T: 'static + DiagnosticDisplay>(
        mut self,
        message: T,
        location: Location,
    ) -> Self {
        self.0
            .related_information
            .push(DiagnosticRelatedInformation {
                message: Box::new(message),
                location,
            });
        self
    }

    /// Some locations, e.g. schema locations where the schema originally comes from
    /// multiple files, have only generated locations. This leads to unhelpful error
    /// annotations. Add a method that skips annotating the error if the location is
    /// generated, so that we can have good (annotated) error messages (if possible)
    /// and sparse (but not unhelpful) error messages when not possible.
    pub fn annotate_if_location_exists<T: 'static + DiagnosticDisplay>(
        self,
        message: T,
        location: Location,
    ) -> Self {
        if !location.source_location().is_generated() {
            self.annotate(message, location)
        } else {
            self
        }
    }

    pub fn metadata_for_machine(mut self, key: impl AsRef<str>, value: impl AsRef<str>) -> Self {
        self.0
            .machine_readable
            .insert(key.as_ref().to_string(), value.as_ref().to_string());
        self
    }

    pub fn message(&self) -> &impl DiagnosticDisplay {
        &self.0.message
    }

    pub fn location(&self) -> Location {
        self.0.location
    }

    pub fn get_data(&self) -> &[impl DiagnosticDisplay] {
        &self.0.data
    }

    pub fn severity(&self) -> DiagnosticSeverity {
        self.0.severity
    }

    pub fn tags(&self) -> Vec<DiagnosticTag> {
        self.0.tags.clone()
    }

    pub fn machine_readable(&self) -> Option<BTreeMap<String, String>> {
        if self.0.machine_readable.is_empty() {
            None
        } else {
            Some(self.0.machine_readable.clone())
        }
    }

    /// Override the location. This should only be used for exceptional situations.
    /// Typically, diagnostics should be constructed with a correct location.
    pub fn override_location(&mut self, location: Location) {
        assert!(
            self.0.location.source_location() == SourceLocationKey::Generated,
            "Diagnostic::override_location can only be called when the location is generated."
        );
        self.0.location = location;
    }

    /// Override the severity. This should only be used for escalating
    /// diagnostics for error reporting. For example, any warnings that
    /// need to be reported as errors can be reconstructed as diagnostics
    /// with a severity of DiagnosticSeverity::ERROR.
    pub fn override_severity(&mut self, severity: DiagnosticSeverity) {
        assert!(
            self.0.severity >= severity, // NOTE: The most critical severity level is actually the lowest enum value
            "Diagnostic::override_severity can only be called when increasing the severity level",
        );

        self.0.severity = severity;
    }

    pub fn related_information(&self) -> &[DiagnosticRelatedInformation] {
        &self.0.related_information
    }

    pub fn print_without_source(&self) -> String {
        let mut result = String::new();
        writeln!(
            result,
            "{message}:{location:?}",
            message = &self.0.message,
            location = self.0.location
        )
        .unwrap();
        if !self.0.related_information.is_empty() {
            for (ix, related) in self.0.related_information.iter().enumerate() {
                writeln!(
                    result,
                    "[related {ix}] {message}:{location:?}",
                    ix = ix + 1,
                    message = related.message,
                    location = related.location
                )
                .unwrap();
            }
        };
        result
    }
}

impl fmt::Display for Diagnostic {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        write!(f, "{}", self.print_without_source())
    }
}

impl Error for Diagnostic {}

// statically verify that the Diagnostic type is thread safe
fn _assert_diagnostic_constraints()
where
    Diagnostic: Send + Sync,
{
}

#[derive(fmt::Debug)]
struct DiagnosticData {
    /// Human readable error message.
    message: Box<dyn DiagnosticDisplay>,

    /// The primary location of this diagnostic.
    location: Location,

    /// Related diagnostic information, such as other definitions in the case of
    /// a duplicate definition error.
    related_information: Vec<DiagnosticRelatedInformation>,

    tags: Vec<DiagnosticTag>,

    severity: DiagnosticSeverity,

    /// A list with data that can be passed to the code actions
    /// `data` is used in the LSP protocol:
    /// @see https://microsoft.github.io/language-server-protocol/specifications/specification-current/#diagnostic
    data: Vec<Box<dyn DiagnosticDisplay>>,

    /// Metadata with (K,V) are strings that can read by machine
    machine_readable: BTreeMap<String, String>,
}

/// Secondary locations attached to a diagnostic.
#[derive(fmt::Debug)]
pub struct DiagnosticRelatedInformation {
    /// The message of this related diagnostic information.
    pub message: Box<dyn DiagnosticDisplay>,

    /// The location of this related diagnostic information.
    pub location: Location,
}

pub trait WithDiagnosticData {
    fn get_data(&self) -> Vec<Box<dyn DiagnosticDisplay>>;
}

/// Trait for diagnostic messages to allow structs that capture
/// some data and can lazily convert it to a message.
pub trait DiagnosticDisplay: fmt::Debug + fmt::Display + Send + Sync {}

/// Automatically implement the trait if constraints are met, so that
/// implementors don't need to.
impl<T> DiagnosticDisplay for T where T: fmt::Debug + fmt::Display + Send + Sync {}

impl From<Diagnostic> for Diagnostics {
    fn from(diagnostic: Diagnostic) -> Self {
        vec![diagnostic]
    }
}

pub fn get_diagnostics_data(diagnostic: &Diagnostic) -> Option<Value> {
    let diagnostic_data = diagnostic.get_data();
    if !diagnostic_data.is_empty() {
        Some(Value::Array(
            diagnostic_data
                .iter()
                .map(|item| Value::String(item.to_string()))
                .collect(),
        ))
    } else {
        None
    }
}


/*
 * Copyright (c) Meta Platforms, Inc. and affiliates.
 *
 * This source code is licensed under the MIT license found in the
 * LICENSE file in the root directory of this source tree.
 */

use crate::DiagnosticSeverity;
use crate::Diagnostics;

#[derive(Debug)]
pub struct CriticalDiagnostics(pub Diagnostics);
#[derive(Debug)]
pub struct StableDiagnostics(pub Diagnostics);

pub type DiagnosticCheck = Result<StableDiagnostics, CriticalDiagnostics>;

/// Checks if any results with diagnostic data are fatal to the compiler
fn check(critical_severity: DiagnosticSeverity, diagnostics: Diagnostics) -> DiagnosticCheck {
    let most_critical_severity = diagnostics
        .iter()
        .map(|diagnostic| diagnostic.severity())
        .min(); // DiagnosticSeverity's lowest value is the highest severity, so we use min()

    if let Some(most_critical_severity) = most_critical_severity {
        if most_critical_severity <= critical_severity {
            return Err(CriticalDiagnostics(diagnostics));
        }
    }

    Ok(StableDiagnostics(diagnostics))
}

/// Escalate all diagnostics at or above the critical severity to the highest severity (DiagnosticSeverity::ERROR)
fn escalate(critical_severity: DiagnosticSeverity, diagnostics: &mut Diagnostics) {
    diagnostics.iter_mut().for_each(|diagnostic| {
        if diagnostic.severity() <= critical_severity {
            diagnostic.override_severity(DiagnosticSeverity::ERROR);
        }
    });
}

/// Convenience function that runs escalate() followed by check()
pub fn escalate_and_check(
    critical_severity: DiagnosticSeverity,
    mut diagnostics: Diagnostics,
) -> DiagnosticCheck {
    escalate(critical_severity, &mut diagnostics);
    check(critical_severity, diagnostics)
}

#[cfg(test)]
mod check_tests {
    use super::*;
    use crate::Diagnostic;
    use crate::Location;

    #[test]
    fn test_empty_diagnostics() {
        let result = check(DiagnosticSeverity::HINT, Vec::new());

        assert!(result.is_ok());
    }

    #[test]
    fn test_only_noncritical_results() {
        let diagnostics = vec![
            Diagnostic::hint("A hint", Location::generated(), Vec::new()),
            Diagnostic::hint("Some information", Location::generated(), Vec::new()),
        ];

        let result = check(DiagnosticSeverity::WARNING, diagnostics);

        assert!(result.is_ok());
    }

    #[test]
    fn test_only_critical_results() {
        let diagnostics = vec![
            Diagnostic::error("An error", Location::generated()),
            Diagnostic::warning("Some warning", Location::generated(), Vec::new()),
        ];

        let result = check(DiagnosticSeverity::WARNING, diagnostics);

        assert!(result.is_err());
    }

    #[test]
    fn test_some_critical_results() {
        let diagnostics = vec![
            Diagnostic::error("An error", Location::generated()),
            Diagnostic::warning("Some warning", Location::generated(), Vec::new()),
            // Only the last item should keep its severity
            Diagnostic::hint("An unfortunate hint", Location::generated(), Vec::new()),
        ];

        let result = check(DiagnosticSeverity::WARNING, diagnostics);

        assert!(result.is_err());
    }
}

#[cfg(test)]
mod escalate_tests {
    use super::*;
    use crate::Diagnostic;
    use crate::Location;

    #[test]
    fn test_empty_diagnostics() {
        let mut diagnostics = Vec::new();
        escalate(DiagnosticSeverity::ERROR, &mut diagnostics);
        assert!(diagnostics.is_empty());
    }

    #[test]
    fn test_no_escalation() {
        let mut diagnostics = vec![
            Diagnostic::hint("A hint", Location::generated(), Vec::new()),
            Diagnostic::hint("Some information", Location::generated(), Vec::new()),
            Diagnostic::hint("More information", Location::generated(), Vec::new()),
        ];

        escalate(DiagnosticSeverity::WARNING, &mut diagnostics);

        assert!(
            diagnostics
                .iter()
                .all(|d| d.severity() == DiagnosticSeverity::HINT)
        );
    }

    #[test]
    fn test_all_escalated() {
        let mut diagnostics = vec![
            Diagnostic::error("An error", Location::generated()),
            Diagnostic::warning("Some warning", Location::generated(), Vec::new()),
            Diagnostic::warning("Another warning", Location::generated(), Vec::new()),
            Diagnostic::warning("Yet another warning", Location::generated(), Vec::new()),
        ];

        escalate(DiagnosticSeverity::WARNING, &mut diagnostics);

        assert!(
            diagnostics
                .iter()
                .all(|d| d.severity() == DiagnosticSeverity::ERROR)
        );
    }

    #[test]
    fn test_some_escalations() {
        let mut diagnostics = vec![
            Diagnostic::error("An error", Location::generated()),
            Diagnostic::warning("Some warning", Location::generated(), Vec::new()),
            // Only the last item should keep its severity
            Diagnostic::hint("An unfortunate hint", Location::generated(), Vec::new()),
        ];

        escalate(DiagnosticSeverity::WARNING, &mut diagnostics);

        let expected_severities = vec![
            DiagnosticSeverity::ERROR,
            DiagnosticSeverity::ERROR,
            DiagnosticSeverity::HINT,
        ];

        assert_eq!(diagnostics.len(), expected_severities.len()); // Sanity check
        for i in 0..diagnostics.len() {
            assert_eq!(diagnostics[i].severity(), expected_severities[i]);
        }
    }
}


/*
 * Copyright (c) Meta Platforms, Inc. and affiliates.
 *
 * This source code is licensed under the MIT license found in the
 * LICENSE file in the root directory of this source tree.
 */

use std::fmt::Display;
use std::fmt::Formatter;
use std::fmt::Result as FmtResult;

use indexmap::IndexSet;
use intern::string_key::StringKey;
use intern::Lookup;
use serde::Deserialize;
use serde::Serialize;

use crate::Rollout;

#[derive(Default, Debug, Serialize, Deserialize, Clone)]
#[serde(deny_unknown_fields)]
pub struct FeatureFlags {
    #[serde(default)]
    pub enable_relay_resolver_transform: bool,

    /// Enable deprecated `@outputType` on Relay Resolvers.
    #[serde(default)]
    pub relay_resolver_enable_output_type: FeatureFlag,

    /// For now, this also disallows fragments with variable definitions
    /// This also makes @module to opt in using @no_inline internally
    /// NOTE that the presence of a fragment in this list only controls whether a fragment is *allowed* to
    /// use @no_inline: whether the fragment is inlined or not depends on whether it actually uses that
    /// directive.
    #[serde(default)]
    pub no_inline: FeatureFlag,

    #[serde(default)]
    pub enable_3d_branch_arg_generation: bool,

    #[serde(default)]
    pub actor_change_support: FeatureFlag,

    /// Enable generation of text artifacts used to generate full query strings
    /// later.
    #[serde(default)]
    pub text_artifacts: FeatureFlag,

    #[serde(default)]
    pub skip_printing_nulls: FeatureFlag,

    /// Enable support for the experimental `@alias` directive on fragment spreads.
    #[serde(default)]
    pub enable_fragment_aliases: FeatureFlag,

    /// Print queries in compact form
    #[serde(default)]
    pub compact_query_text: FeatureFlag,

    /// Create normalization nodes for client edges to client objects
    #[serde(default)]
    pub emit_normalization_nodes_for_client_edges: bool,
}

#[derive(Debug, Deserialize, Clone, Serialize)]
#[serde(tag = "kind", rename_all = "lowercase")]
pub enum FeatureFlag {
    /// Fully disabled: developers may not use this feature
    Disabled,

    /// Fully enabled: developers may use this feature
    Enabled,

    /// Partially enabled: developers may only use this feature on the listed items (fragments, fields, types).
    Limited { allowlist: IndexSet<StringKey> },

    /// Partially enabled: used for gradual rollout of the feature
    Rollout { rollout: Rollout },
}

impl Default for FeatureFlag {
    fn default() -> Self {
        FeatureFlag::Disabled
    }
}

impl FeatureFlag {
    pub fn is_enabled_for(&self, name: StringKey) -> bool {
        match self {
            FeatureFlag::Enabled => true,
            FeatureFlag::Limited { allowlist } => allowlist.contains(&name),
            FeatureFlag::Rollout { rollout } => rollout.check(name.lookup()),
            FeatureFlag::Disabled => false,
        }
    }

    pub fn is_fully_enabled(&self) -> bool {
        matches!(self, FeatureFlag::Enabled)
    }
}

impl Display for FeatureFlag {
    fn fmt(&self, f: &mut Formatter<'_>) -> FmtResult {
        match self {
            FeatureFlag::Disabled => f.write_str("disabled"),
            FeatureFlag::Enabled => f.write_str("enabled"),
            FeatureFlag::Limited { allowlist } => {
                let items: Vec<_> = allowlist.iter().map(|x| x.lookup()).collect();
                f.write_str("limited to: ")?;
                f.write_str(&items.join(", "))
            }
            FeatureFlag::Rollout { rollout } => write!(f, "Rollout: {:#?}", rollout),
        }
    }
}


/*
 * Copyright (c) Meta Platforms, Inc. and affiliates.
 *
 * This source code is licensed under the MIT license found in the
 * LICENSE file in the root directory of this source tree.
 */

#![deny(warnings)]
#![deny(rust_2018_idioms)]
#![deny(clippy::all)]

mod console_logger;
mod diagnostic;
mod diagnostic_check;
mod feature_flags;
mod location;
mod named_item;
mod perf_logger;
mod pointer_address;
mod rollout;
mod span;
pub mod sync;
mod text_source;

pub use console_logger::print_time;
pub use console_logger::ConsoleLogEvent;
pub use console_logger::ConsoleLogger;
pub use diagnostic::diagnostics_result;
pub use diagnostic::get_diagnostics_data;
pub use diagnostic::Diagnostic;
pub use diagnostic::DiagnosticDisplay;
pub use diagnostic::DiagnosticRelatedInformation;
pub(crate) use diagnostic::Diagnostics;
pub use diagnostic::DiagnosticsResult;
pub use diagnostic::WithDiagnosticData;
pub use diagnostic::WithDiagnostics;
pub use diagnostic_check::escalate_and_check;
pub use diagnostic_check::CriticalDiagnostics;
pub use diagnostic_check::DiagnosticCheck;
pub use diagnostic_check::StableDiagnostics;
pub use feature_flags::FeatureFlag;
pub use feature_flags::FeatureFlags;
pub use location::Location;
pub use location::SourceLocationKey;
pub use location::WithLocation;
pub use lsp_types::DiagnosticSeverity;
pub use lsp_types::DiagnosticTag;
pub use named_item::ArgumentName;
pub use named_item::DirectiveName;
pub use named_item::EnumName;
pub use named_item::InputObjectName;
pub use named_item::InterfaceName;
pub use named_item::Named;
pub use named_item::NamedItem;
pub use named_item::ObjectName;
pub use named_item::ScalarName;
pub use perf_logger::NoopPerfLogger;
pub use perf_logger::NoopPerfLoggerEvent;
pub use perf_logger::PerfLogEvent;
pub use perf_logger::PerfLogger;
pub use pointer_address::PointerAddress;
pub use rollout::Rollout;
pub use span::Span;
pub use text_source::TextSource;


/*
 * Copyright (c) Meta Platforms, Inc. and affiliates.
 *
 * This source code is licensed under the MIT license found in the
 * LICENSE file in the root directory of this source tree.
 */

use core::cmp::Ordering;
use std::fmt;
use std::path::PathBuf;

use intern::string_key::Intern;
use intern::string_key::StringKey;
use intern::Lookup;

use crate::span::Span;

/// The location of a source. Could be a standalone file (e.g. test.graphql),
/// an embedded source (GraphQL tag in a JS file) or generated code without a
/// location.
#[derive(Copy, Clone, Debug, Eq, PartialEq, Ord, PartialOrd, Hash)]
pub enum SourceLocationKey {
    /// A source embedded within a file. The 0-based index is an index into the
    /// embedded sources. E.g. the second graphql tag has index 1.
    Embedded {
        path: StringKey,
        index: u16,
    },
    Standalone {
        path: StringKey,
    },
    Generated,
}

impl SourceLocationKey {
    /// Returns a `SourceLocationKey` that's not backed by a real file. In most
    /// cases it's preferred to use a related real file.
    pub fn generated() -> Self {
        SourceLocationKey::Generated
    }

    pub fn standalone(path: &str) -> Self {
        SourceLocationKey::Standalone {
            path: path.intern(),
        }
    }
    pub fn embedded(path: &str, index: usize) -> Self {
        SourceLocationKey::Embedded {
            path: path.intern(),
            index: index.try_into().unwrap(),
        }
    }

    pub fn path(self) -> &'static str {
        match self {
            SourceLocationKey::Embedded { path, .. } => path.lookup(),
            SourceLocationKey::Standalone { path } => path.lookup(),
            SourceLocationKey::Generated => "<generated>",
        }
    }

    pub fn get_dir(self) -> PathBuf {
        let mut path = PathBuf::from(self.path());
        path.pop();
        path
    }

    pub fn is_generated(&self) -> bool {
        matches!(self, SourceLocationKey::Generated)
    }
}

/// An absolute source location describing both the file and position (span)
/// with that file.
#[derive(Copy, Clone, Eq, PartialEq, Ord, PartialOrd, Hash)]
pub struct Location {
    /// The source containing this location (e.g. embedded or standalone file).
    source_location: SourceLocationKey,

    /// Relative position with the file
    span: Span,
}

impl fmt::Debug for Location {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        write!(f, "{}:{:?}", self.source_location.path(), self.span)
    }
}

impl Location {
    /// Returns a location that's not backed by a real file. In most cases it's
    /// preferred to use a related real location.
    pub fn generated() -> Self {
        Location::new(SourceLocationKey::generated(), Span::new(0, 0))
    }

    pub fn new(source_location: SourceLocationKey, span: Span) -> Self {
        Self {
            source_location,
            span,
        }
    }

    pub fn source_location(&self) -> SourceLocationKey {
        self.source_location
    }

    pub fn span(&self) -> Span {
        self.span
    }

    pub fn with_span(&self, span: Span) -> Self {
        Self {
            source_location: self.source_location,
            span,
        }
    }

    pub fn contains(&self, subspan: Span) -> bool {
        self.span.contains(subspan)
    }
}

#[derive(Copy, Clone, Debug, Hash, PartialEq, Eq)]
pub struct WithLocation<T> {
    pub location: Location,
    pub item: T,
}

impl<T: Ord> Ord for WithLocation<T> {
    fn cmp(&self, other: &Self) -> Ordering {
        self.item.cmp(&other.item)
    }
}

impl<T: PartialOrd> PartialOrd for WithLocation<T> {
    fn partial_cmp(&self, other: &Self) -> Option<Ordering> {
        self.item.partial_cmp(&other.item)
    }
}

impl<T> WithLocation<T> {
    pub fn from_span(source_location: SourceLocationKey, span: Span, item: T) -> Self {
        Self {
            location: Location::new(source_location, span),
            item,
        }
    }

    pub fn new(location: Location, item: T) -> Self {
        Self { location, item }
    }

    /// Wraps the given item in a `WithLocation`, without associating it with
    /// any real location. In most cases it's preferred to use a related real
    /// location the item was derived from.
    pub fn generated(item: T) -> Self {
        Self {
            location: Location::generated(),
            item,
        }
    }

    pub fn map<U>(self, f: impl FnOnce(T) -> U) -> WithLocation<U> {
        WithLocation {
            location: self.location,
            item: f(self.item),
        }
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    #[test]
    fn with_location_memory() {
        assert_eq!(20, std::mem::size_of::<WithLocation<StringKey>>());
    }

    #[test]
    fn location_debug_printing() {
        assert_eq!(
            format!(
                "{:?}",
                Location::new(
                    SourceLocationKey::embedded("example/file.js", 2),
                    Span::new(10, 30)
                )
            ),
            "example/file.js:10:30".to_string()
        );
    }
}


/*
 * Copyright (c) Meta Platforms, Inc. and affiliates.
 *
 * This source code is licensed under the MIT license found in the
 * LICENSE file in the root directory of this source tree.
 */

use std::fmt;
use std::str::FromStr;

use intern::impl_lookup;
use intern::string_key::Intern;
use intern::string_key::StringKey;
use intern::Lookup;
use serde::Deserialize;
use serde::Serialize;

/// Utility to access an item in a list by its name where name is defined by the
/// `Named` trait.
pub trait NamedItem<'a, T: Named> {
    fn named(self, name: <T as Named>::Name) -> Option<&'a T>;
}

impl<'a, T: Named> NamedItem<'a, T> for &'a [T] {
    fn named(self, name: <T as Named>::Name) -> Option<&'a T> {
        self.iter().find(|x| x.name() == name)
    }
}

/// Represents a node that has a name such as an `Argument` or `Directive`.
/// The Name associated type should be a newtype wrapper around StringKey.
pub trait Named {
    type Name: Eq + PartialEq<<Self as Named>::Name> + Lookup;
    fn name(&self) -> Self::Name;
}

/// Wrapper struct for clarity rather than having StringKey everywhere.
#[derive(
    Clone,
    Copy,
    Debug,
    Deserialize,
    Eq,
    Hash,
    Ord,
    PartialEq,
    PartialOrd,
    Serialize
)]
pub struct DirectiveName(pub StringKey);

impl fmt::Display for DirectiveName {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        fmt::Display::fmt(&self.0, f)
    }
}

impl FromStr for DirectiveName {
    type Err = std::convert::Infallible;
    fn from_str(s: &str) -> Result<Self, Self::Err> {
        Ok(DirectiveName(s.intern()))
    }
}

impl_lookup!(DirectiveName);
#[derive(Clone, Copy, Debug, Eq, PartialEq, Ord, PartialOrd, Hash)]
pub struct ArgumentName(pub StringKey);

impl fmt::Display for ArgumentName {
    fn fmt(&self, fmt: &mut fmt::Formatter<'_>) -> fmt::Result {
        write!(fmt, "{}", self.0)
    }
}

#[derive(
    Clone,
    Copy,
    Debug,
    Eq,
    PartialEq,
    Ord,
    PartialOrd,
    Hash,
    Serialize,
    Deserialize
)]
pub struct ScalarName(pub StringKey);

impl_lookup!(ScalarName);

impl fmt::Display for ScalarName {
    fn fmt(&self, fmt: &mut fmt::Formatter<'_>) -> fmt::Result {
        write!(fmt, "{}", self.0)
    }
}
impl_lookup!(ArgumentName);
#[derive(Clone, Copy, Debug, Eq, PartialEq, Ord, PartialOrd, Hash)]
pub struct ObjectName(pub StringKey);

impl fmt::Display for ObjectName {
    fn fmt(&self, fmt: &mut fmt::Formatter<'_>) -> fmt::Result {
        write!(fmt, "{}", self.0)
    }
}

impl_lookup!(ObjectName);

#[derive(Clone, Copy, Debug, Eq, PartialEq, Ord, PartialOrd, Hash)]
pub struct InputObjectName(pub StringKey);

impl fmt::Display for InputObjectName {
    fn fmt(&self, fmt: &mut fmt::Formatter<'_>) -> fmt::Result {
        write!(fmt, "{}", self.0)
    }
}

impl_lookup!(InputObjectName);

#[derive(Clone, Copy, Debug, Eq, PartialEq, Ord, PartialOrd, Hash)]
pub struct EnumName(pub StringKey);

impl fmt::Display for EnumName {
    fn fmt(&self, fmt: &mut fmt::Formatter<'_>) -> fmt::Result {
        write!(fmt, "{}", self.0)
    }
}
impl_lookup!(EnumName);

#[derive(Clone, Copy, Debug, Eq, PartialEq, Ord, PartialOrd, Hash)]
pub struct InterfaceName(pub StringKey);

impl fmt::Display for InterfaceName {
    fn fmt(&self, fmt: &mut fmt::Formatter<'_>) -> fmt::Result {
        write!(fmt, "{}", self.0)
    }
}

impl_lookup!(InterfaceName);


/*
 * Copyright (c) Meta Platforms, Inc. and affiliates.
 *
 * This source code is licensed under the MIT license found in the
 * LICENSE file in the root directory of this source tree.
 */

pub trait PerfLogger: Send + Sync {
    type PerfLogEvent: PerfLogEvent + Send;
    /// Create log event
    fn create_event(&self, name: &'static str) -> Self::PerfLogEvent;
}

pub trait PerfLogEvent: Send + Sync {
    type Timer: Send + Sync;

    /// Log number
    fn number(&self, name: &'static str, number: usize);

    /// Provides a possibility to log additional fields describing current run (like, project name)
    fn string(&self, name: &'static str, value: String);

    /// Start new execution timer with the name
    fn start(&self, name: &'static str) -> Self::Timer;

    /// Stop timer and log execution time
    fn stop(&self, timer: Self::Timer);

    /// Measure a time of calling a callback.
    fn time<T, F>(&self, name: &'static str, f: F) -> T
    where
        F: FnOnce() -> T,
    {
        let timer = self.start(name);
        let res = f();
        self.stop(timer);
        res
    }

    /// Log the event
    fn complete(self);
}

pub struct NoopPerfLogger;
impl PerfLogger for NoopPerfLogger {
    type PerfLogEvent = NoopPerfLoggerEvent;
    fn create_event(&self, _name: &'static str) -> Self::PerfLogEvent {
        NoopPerfLoggerEvent
    }
}

pub struct NoopPerfLoggerEvent;
impl PerfLogEvent for NoopPerfLoggerEvent {
    type Timer = ();
    fn number(&self, _name: &'static str, _number: usize) {}
    fn string(&self, _name: &'static str, _value: String) {}
    fn start(&self, _name: &'static str) -> Self::Timer {}
    fn stop(&self, _timer: Self::Timer) {}
    fn complete(self) {}
}


/*
 * Copyright (c) Meta Platforms, Inc. and affiliates.
 *
 * This source code is licensed under the MIT license found in the
 * LICENSE file in the root directory of this source tree.
 */

/// A wrapper type that allows comparing pointer equality of references. Two
/// `PointerAddress` values are equal if they point to the same memory location.
///
/// This type is _sound_, but misuse can easily lead to logical bugs if the memory
/// of one PointerAddress could not have been freed and reused for a subsequent
/// PointerAddress.
#[derive(Hash, Eq, PartialEq, Clone, Copy)]
pub struct PointerAddress(usize);

impl PointerAddress {
    pub fn new<T>(ptr: &T) -> Self {
        let ptr_address: usize = unsafe { std::mem::transmute(ptr) };
        Self(ptr_address)
    }
}


/*
 * Copyright (c) Meta Platforms, Inc. and affiliates.
 *
 * This source code is licensed under the MIT license found in the
 * LICENSE file in the root directory of this source tree.
 */

use md5::Digest;
use md5::Md5;
use serde::Deserialize;
use serde::Serialize;

/// A utility to enable gradual rollout of large codegen changes.
/// Can be constructed as the Default which passes or a percentage between 0 and
/// 100.
#[derive(Default, Debug, Serialize, Deserialize, Clone, Copy)]
pub struct Rollout(Option<u8>);

impl Rollout {
    /// Checks some key deterministically and passes on average the given
    /// percentage of the rollout.
    /// A typical key to pass in could be the fragment or operation name.
    pub fn check(&self, key: impl AsRef<[u8]>) -> bool {
        if let Some(percentage) = self.0 {
            let hash = Md5::digest(key.as_ref());
            let hash: u16 = ((hash[1] as u16) << 8) | (hash[0] as u16);
            (hash % 100) < (percentage as u16)
        } else {
            true
        }
    }
}


/*
 * Copyright (c) Meta Platforms, Inc. and affiliates.
 *
 * This source code is licensed under the MIT license found in the
 * LICENSE file in the root directory of this source tree.
 */

use std::fmt;

#[derive(Copy, Clone, Eq, PartialEq, Ord, PartialOrd, Hash)]
pub struct Span {
    pub start: u32,
    pub end: u32,
}

impl Span {
    pub fn empty() -> Self {
        Self { start: 0, end: 0 }
    }

    pub fn new(start: u32, end: u32) -> Self {
        debug_assert!(start <= end);
        Self { start, end }
    }

    pub fn from_usize(start: usize, end: usize) -> Self {
        Self::new(u32::try_from(start).unwrap(), u32::try_from(end).unwrap())
    }

    pub fn as_usize(self) -> (usize, usize) {
        (self.start as usize, self.end as usize)
    }

    // clippy suggest to use `subspan.start < self.start` :-)
    #[allow(clippy::suspicious_operation_groupings)]
    pub fn contains(self, subspan: Span) -> bool {
        subspan.start >= self.start && subspan.start < self.end && subspan.end <= self.end
    }

    pub fn with_offset(self, offset: u32) -> Self {
        Self::new(self.start + offset, self.end + offset)
    }
}

impl fmt::Debug for Span {
    fn fmt(&self, fmt: &mut fmt::Formatter<'_>) -> fmt::Result {
        write!(fmt, "{}:{}", self.start, self.end)
    }
}

impl From<std::ops::Range<usize>> for Span {
    fn from(range: std::ops::Range<usize>) -> Self {
        Span::from_usize(range.start, range.end)
    }
}

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

    #[test]
    fn span_contains() {
        let outer_span = Span::new(1, 10);
        let overflow_right_start_span = Span::new(10, 10);
        let overflow_right_end_span = Span::new(1, 11);
        let overflow_left_span_start = Span::new(0, 10);

        assert!(outer_span.contains(outer_span), "A span contains itself");
        assert!(
            !outer_span.contains(overflow_right_start_span),
            "A span doesn't contain a subspan whose start is equal to the end of the current span"
        );
        assert!(
            !outer_span.contains(overflow_right_end_span),
            "A span doesn't contain a subspan whose end is greater than the end of the current span"
        );
        assert!(
            !outer_span.contains(overflow_left_span_start),
            "A span doesn't contain a subspan whose start is less than the start of the current span",
        );
    }
}


/*
 * Copyright (c) Meta Platforms, Inc. and affiliates.
 *
 * This source code is licensed under the MIT license found in the
 * LICENSE file in the root directory of this source tree.
 */

pub use std::sync::atomic::Ordering;
pub use std::sync::atomic::Ordering::SeqCst;

pub use rayon::prelude::IntoParallelIterator;
pub use rayon::prelude::ParallelIterator;

/// An abstraction over rayon's par_iter that falls back to a normal
/// iterator when compiling to WebAssembly.
#[cfg(target_arch = "wasm32")]
pub fn par_iter<T: IntoIterator + rayon::iter::IntoParallelIterator>(t: T) -> T::IntoIter {
    t.into_iter()
}

/// An abstraction over rayon's par_iter that falls back to a normal
/// iterator when compiling to WebAssembly.
#[cfg(not(target_arch = "wasm32"))]
pub fn par_iter<T: IntoIterator + rayon::iter::IntoParallelIterator>(t: T) -> T::Iter {
    t.into_par_iter()
}

#[cfg(not(target_arch = "wasm32"))]
pub fn try_join<T1, F1, T2, F2, E>(f1: F1, f2: F2) -> Result<(T1, T2), Vec<E>>
where
    F1: FnOnce() -> Result<T1, Vec<E>> + Send,
    F2: FnOnce() -> Result<T2, Vec<E>> + Send,
    T1: Send,
    T2: Send,
    E: Send,
{
    let (v1, v2) = rayon::join(f1, f2);
    Ok((v1?, v2?))
}

#[cfg(target_arch = "wasm32")]
pub fn try_join<T1, F1, T2, F2, E>(f1: F1, f2: F2) -> Result<(T1, T2), Vec<E>>
where
    F1: FnOnce() -> Result<T1, Vec<E>> + Send,
    F2: FnOnce() -> Result<T2, Vec<E>> + Send,
    T1: Send,
    T2: Send,
    E: Send,
{
    Ok((f1()?, f2()?))
}