1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106

use std::collections::{HashMap, HashSet};
use std::thread;

use itertools::*;
use rayon::prelude::*;

use crate::markoff::Disjoint;
use crate::numbers::{FpNum, GroupElem};

/// Configures tests to be run on orbits of the Markoff graph modulo `P`.
pub struct OrbitTester<const P: u128> {
    targets: HashSet<u128>,
}

/// The results of a successfully run `OrbitTester`.
pub struct OrbitTesterResults {
    results: HashMap<u128, Disjoint<u128>>,
}

type Msg = (u128, u128, u128);

impl<const P: u128> OrbitTester<P> {
    /// Consume and run this `OrbitTester`, blocking until completion, and returning the results.
    /// This method may spawn multiple worker threads, which are guarenteed to be joined before
    /// `run` returns.
    pub fn run(self) -> OrbitTesterResults {
        let mut results = HashMap::with_capacity(self.targets.len());
        for x in &self.targets {
            results.insert(*x, Disjoint::new());
        }

        let mut inv2 = FpNum::<P>::from(2);
        inv2 = inv2.inverse();

        let (tx, rx) = std::sync::mpsc::sync_channel::<Msg>(1024);

        let handle = thread::spawn(move || {
            for (x, y, z) in rx.iter() {
                if results.contains_key(&z) {
                    if let Some(disjoint) = results.get_mut(&x) {
                        disjoint.associate(y, y);
                    }
                    if let Some(disjoint) = results.get_mut(&y) {
                        disjoint.associate(x, z);
                    }
                }
            }

            results
        });

        self.targets
            .iter()
            .combinations_with_replacement(2)
            .map(|v| (v[0], v[1]))
            .par_bridge()
            .for_each(|(x, y)| {
                let x = FpNum::from(*x);
                let y = FpNum::from(*y);

                // We use the non-normalized equation: x^2 + y^2 + z^2 - xyz = 0
                let disc = x * y - 4 * (x * x + y * y);
                let neg_b = x * y;

                match disc.int_sqrt().map(u128::from) {
                    Some(0) => {
                        let z = neg_b * inv2;
                        _ = tx.send((u128::from(x), u128::from(y), u128::from(z)));
                    }
                    Some(root_disc) => {
                        let z = (neg_b + FpNum::from(root_disc)) * inv2;
                        _ = tx.send((u128::from(x), u128::from(y), u128::from(z)));
                        let z = (neg_b - FpNum::from(root_disc)) * inv2;
                        _ = tx.send((u128::from(x), u128::from(y), u128::from(z)));
                    }
                    None => {}
                }
            });
        drop(tx);

        let results = handle.join().unwrap();

        OrbitTesterResults { results }
    }

    /// Creates a new `OrbetTester` with default settings and no targets.
    pub fn new() -> OrbitTester<P> {
        OrbitTester {
            targets: HashSet::new(),
        }
    }

    /// Adds a target order to the list of orders to be tested.
    pub fn add_target(&mut self, t: u128) {
        self.targets.insert(t);
    }
}

impl OrbitTesterResults {
    /// The results of the test, as an iterator yielding each coordinate of a target order, along
    /// with the partitioning of the target orders into disjoint sets, which are subsets of the
    /// orbits under the fixed first coordinate.
    pub fn results(&self) -> impl Iterator<Item = (&u128, &Disjoint<u128>)> {
        self.results.iter()
    }
}