/*
 * 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 dual-licensed under either the MIT license found in the
// LICENSE-MIT file in the root directory of this source tree or the Apache
// License, Version 2.0 found in the LICENSE-APACHE file in the root directory
// of this source tree. You may select, at your option, one of the above-listed
// licenses.

//! Defines the CipherSuite trait to specify the underlying primitives for VOPRF

use digest::core_api::BlockSizeUser;
use digest::{FixedOutput, HashMarker, OutputSizeUser};
use elliptic_curve::VoprfParameters;
use generic_array::typenum::{IsLess, IsLessOrEqual, U256};

use crate::Group;

/// Configures the underlying primitives used in VOPRF
pub trait CipherSuite
where
    <Self::Hash as OutputSizeUser>::OutputSize:
        IsLess<U256> + IsLessOrEqual<<Self::Hash as BlockSizeUser>::BlockSize>,
{
    /// The ciphersuite identifier as dictated by
    /// <https://datatracker.ietf.org/doc/draft-irtf-cfrg-voprf/>
    const ID: &'static str;

    /// A finite cyclic group along with a point representation that allows some
    /// customization on how to hash an input to a curve point. See [`Group`].
    type Group: Group;

    /// The main hash function to use (for HKDF computations and hashing
    /// transcripts).
    type Hash: BlockSizeUser + Default + FixedOutput + HashMarker;
}

impl<T: VoprfParameters> CipherSuite for T
where
    T: Group,
    T::Hash: BlockSizeUser + Default + FixedOutput + HashMarker,
    <T::Hash as OutputSizeUser>::OutputSize:
        IsLess<U256> + IsLessOrEqual<<T::Hash as BlockSizeUser>::BlockSize>,
{
    const ID: &'static str = T::ID;

    type Group = T;

    type Hash = T::Hash;
}


// Copyright (c) Meta Platforms, Inc. and affiliates.
//
// This source code is dual-licensed under either the MIT license found in the
// LICENSE-MIT file in the root directory of this source tree or the Apache
// License, Version 2.0 found in the LICENSE-APACHE file in the root directory
// of this source tree. You may select, at your option, one of the above-listed
// licenses.

//! Common functionality between multiple OPRF modes.

use core::convert::TryFrom;
use core::ops::Add;

use derive_where::derive_where;
use digest::core_api::BlockSizeUser;
use digest::{Digest, Output, OutputSizeUser};
use generic_array::sequence::Concat;
use generic_array::typenum::{IsLess, IsLessOrEqual, Unsigned, U2, U256, U9};
use generic_array::{ArrayLength, GenericArray};
use rand_core::{CryptoRng, RngCore};
use subtle::ConstantTimeEq;

#[cfg(feature = "serde")]
use crate::serialization::serde::{Element, Scalar};
use crate::{CipherSuite, Error, Group, InternalError, Result};

///////////////
// Constants //
// ========= //
///////////////

pub(crate) const STR_FINALIZE: [u8; 8] = *b"Finalize";
pub(crate) const STR_SEED: [u8; 5] = *b"Seed-";
pub(crate) const STR_DERIVE_KEYPAIR: [u8; 13] = *b"DeriveKeyPair";
pub(crate) const STR_COMPOSITE: [u8; 9] = *b"Composite";
pub(crate) const STR_CHALLENGE: [u8; 9] = *b"Challenge";
pub(crate) const STR_INFO: [u8; 4] = *b"Info";
pub(crate) const STR_OPRF: [u8; 7] = *b"OPRFV1-";
pub(crate) const STR_HASH_TO_SCALAR: [u8; 13] = *b"HashToScalar-";
pub(crate) const STR_HASH_TO_GROUP: [u8; 12] = *b"HashToGroup-";

/// Determines the mode of operation (either base mode or verifiable mode). This
/// is only used for custom implementations for [`Group`].
#[derive(Clone, Copy, Debug)]
pub enum Mode {
    /// Non-verifiable mode.
    Oprf,
    /// Verifiable mode.
    Voprf,
    /// Partially-oblivious mode.
    Poprf,
}

impl Mode {
    /// Mode as it is represented in a context string.
    pub fn to_u8(self) -> u8 {
        match self {
            Mode::Oprf => 0,
            Mode::Voprf => 1,
            Mode::Poprf => 2,
        }
    }
}

////////////////////////////
// High-level API Structs //
// ====================== //
////////////////////////////

/// The first client message sent from a client (either verifiable or not) to a
/// server (either verifiable or not).
#[derive_where(Clone, ZeroizeOnDrop)]
#[derive_where(Debug, Eq, Hash, Ord, PartialEq, PartialOrd; <CS::Group as Group>::Elem)]
#[cfg_attr(
    feature = "serde",
    derive(serde::Deserialize, serde::Serialize),
    serde(bound = "")
)]
pub struct BlindedElement<CS: CipherSuite>(
    #[cfg_attr(feature = "serde", serde(with = "Element::<CS::Group>"))]
    pub(crate)  <CS::Group as Group>::Elem,
)
where
    <CS::Hash as OutputSizeUser>::OutputSize:
        IsLess<U256> + IsLessOrEqual<<CS::Hash as BlockSizeUser>::BlockSize>;

/// The server's response to the [BlindedElement] message from a client (either
/// verifiable or not) to a server (either verifiable or not).
#[derive_where(Clone, ZeroizeOnDrop)]
#[derive_where(Debug, Eq, Hash, Ord, PartialEq, PartialOrd; <CS::Group as Group>::Elem)]
#[cfg_attr(
    feature = "serde",
    derive(serde::Deserialize, serde::Serialize),
    serde(bound = "")
)]
pub struct EvaluationElement<CS: CipherSuite>(
    #[cfg_attr(feature = "serde", serde(with = "Element::<CS::Group>"))]
    pub(crate)  <CS::Group as Group>::Elem,
)
where
    <CS::Hash as OutputSizeUser>::OutputSize:
        IsLess<U256> + IsLessOrEqual<<CS::Hash as BlockSizeUser>::BlockSize>;

/// Contains prepared [`EvaluationElement`]s by a server batch evaluate
/// preparation.
#[derive_where(Clone, ZeroizeOnDrop)]
#[derive_where(Debug, Eq, Hash, Ord, PartialEq, PartialOrd; <CS::Group as Group>::Elem)]
#[cfg_attr(
    feature = "serde",
    derive(serde::Deserialize, serde::Serialize),
    serde(bound = "")
)]
pub struct PreparedEvaluationElement<CS: CipherSuite>(pub(crate) EvaluationElement<CS>)
where
    <CS::Hash as OutputSizeUser>::OutputSize:
        IsLess<U256> + IsLessOrEqual<<CS::Hash as BlockSizeUser>::BlockSize>;

/// A proof produced by a server that the OPRF output matches against a server
/// public key.
#[derive_where(Clone, ZeroizeOnDrop)]
#[derive_where(Debug, Eq, Hash, Ord, PartialEq, PartialOrd; <CS::Group as Group>::Scalar)]
#[cfg_attr(
    feature = "serde",
    derive(serde::Deserialize, serde::Serialize),
    serde(bound = "")
)]
pub struct Proof<CS: CipherSuite>
where
    <CS::Hash as OutputSizeUser>::OutputSize:
        IsLess<U256> + IsLessOrEqual<<CS::Hash as BlockSizeUser>::BlockSize>,
{
    #[cfg_attr(feature = "serde", serde(with = "Scalar::<CS::Group>"))]
    pub(crate) c_scalar: <CS::Group as Group>::Scalar,
    #[cfg_attr(feature = "serde", serde(with = "Scalar::<CS::Group>"))]
    pub(crate) s_scalar: <CS::Group as Group>::Scalar,
}

/////////////////////
// Proof Functions //
// =============== //
/////////////////////

/// Can only fail with [`Error::Batch`].
#[allow(clippy::many_single_char_names)]
pub(crate) fn generate_proof<CS: CipherSuite, R: RngCore + CryptoRng>(
    rng: &mut R,
    k: <CS::Group as Group>::Scalar,
    a: <CS::Group as Group>::Elem,
    b: <CS::Group as Group>::Elem,
    cs: impl Iterator<Item = <CS::Group as Group>::Elem> + ExactSizeIterator,
    ds: impl Iterator<Item = <CS::Group as Group>::Elem> + ExactSizeIterator,
    mode: Mode,
) -> Result<Proof<CS>>
where
    <CS::Hash as OutputSizeUser>::OutputSize:
        IsLess<U256> + IsLessOrEqual<<CS::Hash as BlockSizeUser>::BlockSize>,
{
    // https://www.ietf.org/archive/id/draft-irtf-cfrg-voprf-11.html#section-2.2.1

    let (m, z) = compute_composites::<CS, _, _>(Some(k), b, cs, ds, mode)?;

    let r = CS::Group::random_scalar(rng);
    let t2 = a * &r;
    let t3 = m * &r;

    // Bm = GG.SerializeElement(B)
    let bm = CS::Group::serialize_elem(b);
    // a0 = GG.SerializeElement(M)
    let a0 = CS::Group::serialize_elem(m);
    // a1 = GG.SerializeElement(Z)
    let a1 = CS::Group::serialize_elem(z);
    // a2 = GG.SerializeElement(t2)
    let a2 = CS::Group::serialize_elem(t2);
    // a3 = GG.SerializeElement(t3)
    let a3 = CS::Group::serialize_elem(t3);

    let elem_len = <CS::Group as Group>::ElemLen::U16.to_be_bytes();

    // h2Input = I2OSP(len(Bm), 2) || Bm ||
    //           I2OSP(len(a0), 2) || a0 ||
    //           I2OSP(len(a1), 2) || a1 ||
    //           I2OSP(len(a2), 2) || a2 ||
    //           I2OSP(len(a3), 2) || a3 ||
    //           "Challenge"
    let h2_input = [
        &elem_len,
        bm.as_slice(),
        &elem_len,
        &a0,
        &elem_len,
        &a1,
        &elem_len,
        &a2,
        &elem_len,
        &a3,
        &STR_CHALLENGE,
    ];

    let dst = Dst::new::<CS, _, _>(STR_HASH_TO_SCALAR, mode);
    // This can't fail, the size of the `input` is known.
    let c_scalar = CS::Group::hash_to_scalar::<CS::Hash>(&h2_input, &dst.as_dst()).unwrap();
    let s_scalar = r - &(c_scalar * &k);

    Ok(Proof { c_scalar, s_scalar })
}

/// Can only fail with [`Error::ProofVerification`] or [`Error::Batch`].
#[allow(clippy::many_single_char_names)]
pub(crate) fn verify_proof<CS: CipherSuite>(
    a: <CS::Group as Group>::Elem,
    b: <CS::Group as Group>::Elem,
    cs: impl Iterator<Item = <CS::Group as Group>::Elem> + ExactSizeIterator,
    ds: impl Iterator<Item = <CS::Group as Group>::Elem> + ExactSizeIterator,
    proof: &Proof<CS>,
    mode: Mode,
) -> Result<()>
where
    <CS::Hash as OutputSizeUser>::OutputSize:
        IsLess<U256> + IsLessOrEqual<<CS::Hash as BlockSizeUser>::BlockSize>,
{
    // https://www.ietf.org/archive/id/draft-irtf-cfrg-voprf-11.html#section-2.2.2
    let (m, z) = compute_composites::<CS, _, _>(None, b, cs, ds, mode)?;
    let t2 = (a * &proof.s_scalar) + &(b * &proof.c_scalar);
    let t3 = (m * &proof.s_scalar) + &(z * &proof.c_scalar);

    // Bm = GG.SerializeElement(B)
    let bm = CS::Group::serialize_elem(b);
    // a0 = GG.SerializeElement(M)
    let a0 = CS::Group::serialize_elem(m);
    // a1 = GG.SerializeElement(Z)
    let a1 = CS::Group::serialize_elem(z);
    // a2 = GG.SerializeElement(t2)
    let a2 = CS::Group::serialize_elem(t2);
    // a3 = GG.SerializeElement(t3)
    let a3 = CS::Group::serialize_elem(t3);

    let elem_len = <CS::Group as Group>::ElemLen::U16.to_be_bytes();

    // h2Input = I2OSP(len(Bm), 2) || Bm ||
    //           I2OSP(len(a0), 2) || a0 ||
    //           I2OSP(len(a1), 2) || a1 ||
    //           I2OSP(len(a2), 2) || a2 ||
    //           I2OSP(len(a3), 2) || a3 ||
    //           "Challenge"
    let h2_input = [
        &elem_len,
        bm.as_slice(),
        &elem_len,
        &a0,
        &elem_len,
        &a1,
        &elem_len,
        &a2,
        &elem_len,
        &a3,
        &STR_CHALLENGE,
    ];

    let dst = Dst::new::<CS, _, _>(STR_HASH_TO_SCALAR, mode);
    // This can't fail, the size of the `input` is known.
    let c = CS::Group::hash_to_scalar::<CS::Hash>(&h2_input, &dst.as_dst()).unwrap();

    match c.ct_eq(&proof.c_scalar).into() {
        true => Ok(()),
        false => Err(Error::ProofVerification),
    }
}

type ComputeCompositesResult<CS> = (
    <<CS as CipherSuite>::Group as Group>::Elem,
    <<CS as CipherSuite>::Group as Group>::Elem,
);

/// Can only fail with [`Error::Batch`].
fn compute_composites<
    CS: CipherSuite,
    IC: Iterator<Item = <CS::Group as Group>::Elem> + ExactSizeIterator,
    ID: Iterator<Item = <CS::Group as Group>::Elem> + ExactSizeIterator,
>(
    k_option: Option<<CS::Group as Group>::Scalar>,
    b: <CS::Group as Group>::Elem,
    c_slice: IC,
    d_slice: ID,
    mode: Mode,
) -> Result<ComputeCompositesResult<CS>>
where
    <CS::Hash as OutputSizeUser>::OutputSize:
        IsLess<U256> + IsLessOrEqual<<CS::Hash as BlockSizeUser>::BlockSize>,
{
    // https://www.ietf.org/archive/id/draft-irtf-cfrg-voprf-11.html#section-2.2.1

    let elem_len = <CS::Group as Group>::ElemLen::U16.to_be_bytes();

    if c_slice.len() != d_slice.len() {
        return Err(Error::Batch);
    }

    let len = u16::try_from(c_slice.len()).map_err(|_| Error::Batch)?;

    // seedDST = "Seed-" || contextString
    let seed_dst = Dst::new::<CS, _, _>(STR_SEED, mode);

    // h1Input = I2OSP(len(Bm), 2) || Bm ||
    //           I2OSP(len(seedDST), 2) || seedDST
    // seed = Hash(h1Input)
    let seed = CS::Hash::new()
        .chain_update(elem_len)
        .chain_update(CS::Group::serialize_elem(b))
        .chain_update(seed_dst.i2osp_2())
        .chain_update_multi(&seed_dst.as_dst())
        .finalize();
    let seed_len = i2osp_2_array(&seed);

    let mut m = CS::Group::identity_elem();
    let mut z = CS::Group::identity_elem();

    for (i, (c, d)) in (0..len).zip(c_slice.zip(d_slice)) {
        // Ci = GG.SerializeElement(Cs[i])
        let ci = CS::Group::serialize_elem(c);
        // Di = GG.SerializeElement(Ds[i])
        let di = CS::Group::serialize_elem(d);
        // h2Input = I2OSP(len(seed), 2) || seed || I2OSP(i, 2) ||
        //           I2OSP(len(Ci), 2) || Ci ||
        //           I2OSP(len(Di), 2) || Di ||
        //           "Composite"
        let h2_input = [
            seed_len.as_slice(),
            &seed,
            &i.to_be_bytes(),
            &elem_len,
            &ci,
            &elem_len,
            &di,
            &STR_COMPOSITE,
        ];

        let dst = Dst::new::<CS, _, _>(STR_HASH_TO_SCALAR, mode);
        // This can't fail, the size of the `input` is known.
        let di = CS::Group::hash_to_scalar::<CS::Hash>(&h2_input, &dst.as_dst()).unwrap();
        m = c * &di + &m;
        z = match k_option {
            Some(_) => z,
            None => d * &di + &z,
        };
    }

    z = match k_option {
        Some(k) => m * &k,
        None => z,
    };

    Ok((m, z))
}

/////////////////////
// Inner Functions //
// =============== //
/////////////////////

/// Can only fail with [`Error::DeriveKeyPair`] and [`Error::Protocol`].
pub(crate) fn derive_key_internal<CS: CipherSuite>(
    seed: &[u8],
    info: &[u8],
    mode: Mode,
) -> Result<<CS::Group as Group>::Scalar, Error>
where
    <CS::Hash as OutputSizeUser>::OutputSize:
        IsLess<U256> + IsLessOrEqual<<CS::Hash as BlockSizeUser>::BlockSize>,
{
    let dst = Dst::new::<CS, _, _>(STR_DERIVE_KEYPAIR, mode);

    let info_len = i2osp_2(info.len()).map_err(|_| Error::DeriveKeyPair)?;

    for counter in 0_u8..=u8::MAX {
        // deriveInput = seed || I2OSP(len(info), 2) || info
        // skS = G.HashToScalar(deriveInput || I2OSP(counter, 1), DST = "DeriveKeyPair"
        // || contextString)
        let sk_s = CS::Group::hash_to_scalar::<CS::Hash>(
            &[seed, &info_len, info, &counter.to_be_bytes()],
            &dst.as_dst(),
        )
        .map_err(|_| Error::DeriveKeyPair)?;

        if !bool::from(CS::Group::is_zero_scalar(sk_s)) {
            return Ok(sk_s);
        }
    }

    Err(Error::Protocol)
}

/// Corresponds to DeriveKeyPair() function from the VOPRF specification.
///
/// # Errors
/// - [`Error::DeriveKeyPair`] if the `input` and `seed` together are longer
///   then `u16::MAX - 3`.
/// - [`Error::Protocol`] if the protocol fails and can't be completed.
#[cfg(feature = "danger")]
pub fn derive_key<CS: CipherSuite>(
    seed: &[u8],
    info: &[u8],
    mode: Mode,
) -> Result<<CS::Group as Group>::Scalar, Error>
where
    <CS::Hash as OutputSizeUser>::OutputSize:
        IsLess<U256> + IsLessOrEqual<<CS::Hash as BlockSizeUser>::BlockSize>,
{
    derive_key_internal::<CS>(seed, info, mode)
}

type DeriveKeypairResult<CS> = (
    <<CS as CipherSuite>::Group as Group>::Scalar,
    <<CS as CipherSuite>::Group as Group>::Elem,
);

/// Can only fail with [`Error::DeriveKeyPair`] and [`Error::Protocol`].
pub(crate) fn derive_keypair<CS: CipherSuite>(
    seed: &[u8],
    info: &[u8],
    mode: Mode,
) -> Result<DeriveKeypairResult<CS>, Error>
where
    <CS::Hash as OutputSizeUser>::OutputSize:
        IsLess<U256> + IsLessOrEqual<<CS::Hash as BlockSizeUser>::BlockSize>,
{
    let sk_s = derive_key_internal::<CS>(seed, info, mode)?;
    let pk_s = CS::Group::base_elem() * &sk_s;

    Ok((sk_s, pk_s))
}

/// Inner function for blind that assumes that the blinding factor has already
/// been chosen, and therefore takes it as input. Does not check if the blinding
/// factor is non-zero.
///
/// Can only fail with [`Error::Input`].
pub(crate) fn deterministic_blind_unchecked<CS: CipherSuite>(
    input: &[u8],
    blind: &<CS::Group as Group>::Scalar,
    mode: Mode,
) -> Result<<CS::Group as Group>::Elem>
where
    <CS::Hash as OutputSizeUser>::OutputSize:
        IsLess<U256> + IsLessOrEqual<<CS::Hash as BlockSizeUser>::BlockSize>,
{
    let hashed_point = hash_to_group::<CS>(input, mode)?;
    Ok(hashed_point * blind)
}

/// Hashes `input` to a point on the curve
pub(crate) fn hash_to_group<CS: CipherSuite>(
    input: &[u8],
    mode: Mode,
) -> Result<<CS::Group as Group>::Elem>
where
    <CS::Hash as OutputSizeUser>::OutputSize:
        IsLess<U256> + IsLessOrEqual<<CS::Hash as BlockSizeUser>::BlockSize>,
{
    let dst = Dst::new::<CS, _, _>(STR_HASH_TO_GROUP, mode);
    CS::Group::hash_to_curve::<CS::Hash>(&[input], &dst.as_dst()).map_err(|_| Error::Input)
}

/// Internal function that finalizes the hash input for OPRF, VOPRF & POPRF.
/// Returned values can only fail with [`Error::Input`].
pub(crate) fn server_evaluate_hash_input<CS: CipherSuite>(
    input: &[u8],
    info: Option<&[u8]>,
    issued_element: GenericArray<u8, <<CS as CipherSuite>::Group as Group>::ElemLen>,
) -> Result<Output<CS::Hash>>
where
    <CS::Hash as OutputSizeUser>::OutputSize:
        IsLess<U256> + IsLessOrEqual<<CS::Hash as BlockSizeUser>::BlockSize>,
{
    // OPRF & VOPRF
    // hashInput = I2OSP(len(input), 2) || input ||
    //             I2OSP(len(issuedElement), 2) || issuedElement ||
    //             "Finalize"
    // return Hash(hashInput)
    //
    // POPRF
    // hashInput = I2OSP(len(input), 2) || input ||
    //             I2OSP(len(info), 2) || info ||
    //             I2OSP(len(issuedElement), 2) || issuedElement ||
    //             "Finalize"

    let mut hash = CS::Hash::new()
        .chain_update(i2osp_2(input.as_ref().len()).map_err(|_| Error::Input)?)
        .chain_update(input.as_ref());
    if let Some(info) = info {
        hash = hash
            .chain_update(i2osp_2(info.as_ref().len()).map_err(|_| Error::Input)?)
            .chain_update(info.as_ref());
    }
    Ok(hash
        .chain_update(i2osp_2(issued_element.as_ref().len()).map_err(|_| Error::Input)?)
        .chain_update(issued_element)
        .chain_update(STR_FINALIZE)
        .finalize())
}

pub(crate) struct Dst<L: ArrayLength<u8>> {
    dst_1: GenericArray<u8, L>,
    dst_2: &'static str,
}

impl<L: ArrayLength<u8>> Dst<L> {
    pub(crate) fn new<CS: CipherSuite, T, TL: ArrayLength<u8>>(par_1: T, mode: Mode) -> Self
    where
        T: Into<GenericArray<u8, TL>>,
        TL: Add<U9, Output = L>,
        <CS::Hash as OutputSizeUser>::OutputSize:
            IsLess<U256> + IsLessOrEqual<<CS::Hash as BlockSizeUser>::BlockSize>,
    {
        let par_1 = par_1.into();
        // Generates the contextString parameter as defined in
        // <https://datatracker.ietf.org/doc/draft-irtf-cfrg-voprf/>
        let par_2 = GenericArray::from(STR_OPRF)
            .concat([mode.to_u8()].into())
            .concat([b'-'].into());

        let dst_1 = par_1.concat(par_2);
        let dst_2 = CS::ID;

        assert!(
            L::USIZE + dst_2.len() <= u16::MAX.into(),
            "constructed DST longer then {}",
            u16::MAX
        );

        Self { dst_1, dst_2 }
    }

    pub(crate) fn as_dst(&self) -> [&[u8]; 2] {
        [&self.dst_1, self.dst_2.as_bytes()]
    }

    pub(crate) fn i2osp_2(&self) -> [u8; 2] {
        u16::try_from(L::USIZE + self.dst_2.len())
            .unwrap()
            .to_be_bytes()
    }
}

trait DigestExt {
    fn chain_update_multi(self, data: &[&[u8]]) -> Self;
}

impl<T> DigestExt for T
where
    T: Digest,
{
    fn chain_update_multi(mut self, datas: &[&[u8]]) -> Self {
        for data in datas {
            self.update(data)
        }

        self
    }
}

///////////////////////
// Utility Functions //
// ================= //
///////////////////////

pub(crate) fn i2osp_2(input: usize) -> Result<[u8; 2], InternalError> {
    u16::try_from(input)
        .map(|input| input.to_be_bytes())
        .map_err(|_| InternalError::I2osp)
}

pub(crate) fn i2osp_2_array<L: ArrayLength<u8> + IsLess<U256>>(
    _: &GenericArray<u8, L>,
) -> GenericArray<u8, U2> {
    L::U16.to_be_bytes().into()
}


// Copyright (c) Meta Platforms, Inc. and affiliates.
//
// This source code is dual-licensed under either the MIT license found in the
// LICENSE-MIT file in the root directory of this source tree or the Apache
// License, Version 2.0 found in the LICENSE-APACHE file in the root directory
// of this source tree. You may select, at your option, one of the above-listed
// licenses.

//! Errors which are produced during an execution of the protocol

use displaydoc::Display;

/// [`Result`](core::result::Result) shorthand that uses [`Error`].
pub type Result<T, E = Error> = core::result::Result<T, E>;

/// Represents an error in the manipulation of internal cryptographic data
#[derive(Clone, Copy, Debug, Display, Eq, Hash, Ord, PartialEq, PartialOrd)]
pub enum Error {
    /// Size of info is longer then [`u16::MAX`].
    Info,
    /// Size of input is empty or longer then [`u16::MAX`].
    Input,
    /// Size of info and seed together are longer then `u16::MAX - 3`.
    DeriveKeyPair,
    /// Failure to deserialize bytes
    Deserialization,
    /// Batched items are more then [`u16::MAX`] or length don't match.
    Batch,
    /// In verifiable mode, occurs when the proof failed to verify
    ProofVerification,
    /// The protocol has failed and can't be completed.
    Protocol,
}

/// Only used to implement [`Group`](crate::Group).
#[derive(Clone, Copy, Debug, Eq, Hash, Ord, PartialEq, PartialOrd)]
pub enum InternalError {
    /// Size of input is empty or longer then [`u16::MAX`].
    Input,
    /// `input` is longer then [`u16::MAX`].
    I2osp,
}

#[cfg(feature = "std")]
impl std::error::Error for Error {}