pll-clock-solver.rs

use cassowary::{Constraint, Solver, Variable, WeightedRelation::*, strength::*};
use std::{
    collections::HashMap,
    sync::{Arc, Mutex, mpsc},
};
use thread_pool::Task;

const MIN_CLOCK: f64 = 165.0;
const CLK48_ERR: f64 = 0.001;
const MIN_PCLK1: f64 = 42.0;
const MIN_PCLK2: f64 = 83.0;
const STOP_FIRST: bool = false;
const PRINT_BEST: Option<usize> = Some(100);

struct SolveTask {
    solutions: mpsc::Sender<Solution>,
    pllm_val: f64,
    p_val: f64,
}

#[derive(Debug)]
struct Solution {
    input_frequency: f64,
    pllm: f64,
    n: f64,
    p: f64,
    sysclk: f64,
    hclk: f64,
    ahb_prescaler: f64,
    q: f64,
    apb1_prescaler: f64,
    apb2_prescaler: f64,
    pclk1: f64,
    pclk2: f64,
    clk48: f64,
}

trait AppxEq {
    fn appx_eq(&self, other: Self) -> bool;
    fn appx_le(&self, other: Self) -> bool;
    fn appx_ge(&self, other: Self) -> bool;
}

impl AppxEq for f64 {
    fn appx_eq(&self, other: Self) -> bool {
        (*self - other).abs() < (1e-12)
    }

    fn appx_le(&self, other: Self) -> bool {
        self.appx_eq(other) || *self <= other
    }

    fn appx_ge(&self, other: Self) -> bool {
        self.appx_eq(other) || *self >= other
    }
}

impl Solution {
    fn with_input(&self, input_frequency: f64) -> Self {
        let sysclk = input_frequency / self.pllm * self.n / self.p;
        let hclk = sysclk / self.ahb_prescaler;
        Self {
            input_frequency,
            pllm: self.pllm,
            n: self.n,
            p: self.p,
            sysclk,
            hclk,
            q: self.q,
            apb1_prescaler: self.apb1_prescaler,
            apb2_prescaler: self.apb2_prescaler,
            ahb_prescaler: self.ahb_prescaler,
            pclk1: hclk / self.apb1_prescaler,
            pclk2: hclk / self.apb2_prescaler,
            clk48: input_frequency / self.pllm * self.n / self.q,
        }
    }

    fn round_input_down(&self, nearest: f64) -> Self {
        self.with_input((self.input_frequency / nearest).floor() * nearest)
    }

    fn round_input_up(&self, nearest: f64) -> Self {
        self.with_input((self.input_frequency / nearest).ceil() * nearest)
    }

    fn satisfied(&self) -> bool {
        self.input_frequency.appx_ge(4.0) && self.input_frequency.appx_le(26.0)
        && (self.input_frequency / self.pllm).appx_ge(0.95) && (self.input_frequency / self.pllm).appx_le(2.1)
        && (self.input_frequency / self.pllm * self.n).appx_ge(100.0) && (self.input_frequency / self.pllm * self.n).appx_le(432.0)
        && self.sysclk.appx_ge(24.0) && self.sysclk.appx_le(180.0)
        && self.hclk.appx_le(180.0) && self.pclk1.appx_le(45.0) && self.pclk2.appx_le(90.0)
        && self.clk48.appx_ge(48.0 * (1.0 - CLK48_ERR)) && self.clk48.appx_le(48.0 * (1.0 + CLK48_ERR))
    }
}

impl Task for SolveTask {
    fn run(self) {
        for n_val in 50usize..=432 {
            for ahb_prescaler_val in [1, 2, 4, 8, 16, 64, 128, 256, 512] {
                for q_val in 2..=15 {
                    for apb1_prescaler_val in [1, 2, 4, 8, 16] {
                        for apb2_prescaler_val in [1, 2, 4, 8, 16] {
                            let mut s = Solver::new();
                            let input_frequency = Variable::new();
                            let pllm = 1.0 / self.pllm_val;
                            let n = n_val as f64;
                            let p = 1.0 / self.p_val;
                            let sysclk = Variable::new();
                            let hclk = Variable::new();
                            let ahb_prescaler = 1.0 / (ahb_prescaler_val as f64);
                            let q = 1.0 / (q_val as f64);
                            let apb1_prescaler = 1.0 / (apb1_prescaler_val as f64);
                            let apb2_prescaler = 1.0 / (apb2_prescaler_val as f64);
                            let pclk1 = Variable::new();
                            let pclk2 = Variable::new();
                            let clk48 = Variable::new();

                            if s.add_constraints(&[
                                input_frequency | GE(REQUIRED) | 4.0f64,
                                input_frequency | LE(REQUIRED) | 26.0f64,
                                hclk | LE(REQUIRED) | 180f64,
                                sysclk | EQ(REQUIRED) | (input_frequency * (pllm * n * p)),
                                sysclk | GE(REQUIRED) | 24.0,
                                sysclk | LE(REQUIRED) | 180.0,
                                hclk | EQ(REQUIRED) | sysclk * ahb_prescaler,
                                input_frequency * (pllm * n * q) | EQ(REQUIRED) | clk48,
                                hclk * apb1_prescaler | EQ(REQUIRED) | pclk1,
                                hclk * apb2_prescaler | EQ(REQUIRED) | pclk2,
                                (input_frequency * pllm) | GE(REQUIRED) | 0.95,
                                (input_frequency * pllm) | LE(REQUIRED) | 2.1,
                                (input_frequency * (pllm * n)) | GE(REQUIRED) | 100.0,
                                (input_frequency * (pllm * n)) | LE(REQUIRED) | 432.0,
                                clk48 | GE(REQUIRED) | (48.0 * (1.0 - CLK48_ERR)),
                                clk48 | LE(REQUIRED) | (48.0 * (1.0 + CLK48_ERR)),
                                pclk1 | LE(REQUIRED) | 45.0,
                                pclk2 | LE(REQUIRED) | 90.0,
                                pclk1 | GE(REQUIRED) | MIN_PCLK1,
                                pclk2 | GE(REQUIRED) | MIN_PCLK2,
                                hclk | GE(REQUIRED) | MIN_CLOCK,
                            ])
                            .is_err()
                            {
                                continue;
                            }

                            if self.solutions
                                .send(Solution {
                                    input_frequency: s.get_value(input_frequency),
                                    pllm: self.pllm_val,
                                    n: n_val as f64,
                                    p: self.p_val,
                                    q: q_val as f64,
                                    ahb_prescaler: ahb_prescaler_val as f64,
                                    sysclk: s.get_value(sysclk),
                                    hclk: s.get_value(hclk),
                                    apb1_prescaler: apb1_prescaler_val as f64,
                                    apb2_prescaler: apb2_prescaler_val as f64,
                                    pclk1: s.get_value(pclk1),
                                    pclk2: s.get_value(pclk2),
                                    clk48: s.get_value(clk48),
                                })
                                .is_err()
                            {
                                return;
                            };
                        }
                    }
                }
            }
        }
    }
}

fn sort_solutions(solutions: &mut Vec<Solution>) {
    solutions.sort_by(|a, b| {
        // Compute the bucket index: each bucket is of size 2.0
        let a_bucket = (a.hclk / 2.0).floor();
        let b_bucket = (b.hclk / 2.0).floor();

        // First sort: descending order of the bucket
        // Second sort: descending order of input_frequency within the same bucket
        b_bucket
            .partial_cmp(&a_bucket)
            .unwrap() // safe unwrap if hclk is valid
            .then(b.input_frequency.partial_cmp(&a.input_frequency).unwrap())
    });
}

fn main() {
    let (mut sender, mut tp) = thread_pool::ThreadPool::fixed_size(
        std::thread::available_parallelism()
            .map(Into::into)
            .unwrap_or(1),
    );
    let (solution_sender, solution_receiver) = mpsc::channel();
    for pllm_val in 2usize..=64 {
        for p_val in (2usize..=8).step_by(2) {
            sender
                .send(SolveTask {
                    solutions: solution_sender.clone(),
                    pllm_val: pllm_val as f64,
                    p_val: p_val as f64,
                })
                .unwrap();
        }
    }

    let mut solutions = Vec::new();

    let mut last_queued = usize::MAX;
    let mut last_solutions = 0;
    while last_queued > 0 {
        while let Ok(solution) = solution_receiver.try_recv() {
            if !solution.satisfied() {
                panic!("solution unsatisfied: {:#?}", solution);
            }

            for nearest in [1.0, 0.1, 0.01] {
                let s = solution.round_input_down(nearest);
                if s.satisfied() { solutions.push(s); }
                let s = solution.round_input_up(nearest);
                if s.satisfied() { solutions.push(s); }
            }            
        }        

        if STOP_FIRST && solutions.len() > 0 {
            println!("found solution (stopping after first): \n{:#?}", solutions.first().unwrap());
            std::process::exit(0);
        }

        let solutions_count = solutions.len();
        let queued = tp.queued();
        if queued != last_queued || solutions_count != last_solutions {
            println!("{} left (found {} solutions)...", queued, solutions_count);
        }
        last_queued = queued;
        last_solutions = solutions_count;

        std::thread::sleep(std::time::Duration::from_millis(250));
    }

    println!("done; found {} solutions", solutions.len());
    if solutions.len() > 0 {
        sort_solutions(&mut solutions);

        println!("----- BEST SOLUTION -----\n");
        println!("{:#?}", solutions.first().unwrap());
        println!("\n----- OTHER SOLUTIONS -----");
        for solution in solutions.iter().skip(1).take(PRINT_BEST.unwrap_or(solutions.len())) {
            println!("{solution:?}");
        }
    }
}