AI Flow: Reverse Engineering

//! The nets in sections a person can page through: ground, each supply rail, each bus or interface, then what is
//! left by the functional block that holds most of its pins, then the rest.
//!
//! Names decide first, because a name a person gave (or a connector's pin name) says what a net is for. A net whose
//! name says nothing (`N012`, `CU_R4_1`) or that shares a stem with no other net goes to the block in blocks.json
//! that holds most of its pins. Nothing here reads a photo or changes a net.
use crate::netconf::nat_cmp;
use serde_json::Value;
use std::collections::{BTreeMap, BTreeSet, HashMap};

#[derive(Clone, Debug, PartialEq)]
pub struct Section { pub kind: &'static str, pub title: String, pub why: String, pub nets: Vec<String> }

const GROUND: [&str; 12] = ["GND", "AGND", "DGND", "PGND", "SGND", "GNDA", "GNDD", "VSS", "VSSA", "GROUND", "EARTH", "0V"];
const SUPPLY: [&str; 22] = ["VIN", "VCC", "VDD", "VBAT", "VBATT", "VBUS", "VSYS", "VREF", "VPP", "VM", "VS", "VOUT", "VDDA", "VDDIO", "VCCIO", "AVDD", "DVDD", "AVCC", "VCCA", "VSUP", "VMOT", "VPWR"];
/// A name with one of these is a signal about a rail, not the rail.
const CONTROL: [&str; 20] = ["EN", "ENABLE", "ON", "OFF", "OVERRIDE", "SENSE", "SNS", "FB", "ADJ", "PG", "PGOOD", "OK", "ADC", "DET", "SW", "BST", "BOOT", "DISABLE", "NDISABLE", "FAULT"];

fn tokens(name: &str) -> Vec<String> { name.split(|c: char| c == '_' || c == '-' || c == ' ' || c == '/').filter(|t| !t.is_empty()).map(|t| t.to_uppercase()).collect() }

/// `3V3`, `5V`, `12V`, `1.8V`, `3V3A`, with or without a sign.
fn is_voltage(t: &str) -> bool {
    let t = t.trim_start_matches(['+', '-']);
    let Some(v) = t.find('V') else { return false };
    let (a, b) = (&t[..v], &t[v + 1..]);
    !a.is_empty() && a.chars().all(|c| c.is_ascii_digit() || c == '.') && a.chars().next().unwrap().is_ascii_digit()
        && b.trim_end_matches(|c: char| c.is_ascii_uppercase()).chars().all(|c| c.is_ascii_digit()) && b.chars().filter(|c| c.is_ascii_uppercase()).count() <= 2
}

/// A name that says nothing about the net: one the program made up.
pub fn is_auto(name: &str) -> bool {
    let u = name.to_uppercase();
    let numbered = |p: &str| u.strip_prefix(p).map(|r| !r.is_empty() && r.chars().all(|c| c.is_ascii_digit())).unwrap_or(false);
    u.is_empty() || u.starts_with("CU_") || u.starts_with("UNK") || u.starts_with("NET-(") || numbered("N") || numbered("NET") || numbered("NET_") || numbered("N$")
}
pub fn is_ground(name: &str) -> bool { let t = tokens(name); t.iter().any(|x| GROUND.contains(&x.as_str())) || t.first().map(|x| x.starts_with("GND")).unwrap_or(false) }
pub fn is_supply(name: &str) -> bool {
    let t = tokens(name);
    if t.is_empty() || t.iter().any(|x| CONTROL.contains(&x.as_str()) || CONTROL.contains(&x.trim_start_matches('N'))) { return false; }
    let rail = |x: &String| SUPPLY.contains(&x.as_str()) || is_voltage(x);
    let signed = name.starts_with(['+', '-']) && name.chars().nth(1).map(|c| c.is_ascii_digit()).unwrap_or(false);
    signed || rail(&t[0]) || rail(&t[t.len() - 1])
}

fn strip_digits(s: &str) -> &str { s.trim_end_matches(|c: char| c.is_ascii_digit()) }

/// Letters that tell the members of one family apart: phases, channels, high and low, the two wires of a pair.
const FAMILIES: [&str; 6] = ["ABCD", "UVW", "XYZ", "HL", "PN", "RST"];

/// `GHA` gives (`GH`, `A`): a stem and the letter that would tell this net from its siblings.
fn lettered(name: &str) -> Option<(String, char)> {
    let t = tokens(name); let bare = strip_digits(t.first()?);
    let last = bare.chars().last()?;
    (bare.len() >= 3 && last.is_ascii_uppercase()).then(|| (bare[..bare.len() - 1].to_string(), last))
}

/// The stems a signal's name could share with others, most specific first, each with how many nets must share it.
/// `family`: whether a stem has members told apart by one letter of a known set.
fn stems(name: &str, family: &dyn Fn(&str) -> bool) -> Vec<(String, usize)> {
    let t = tokens(name);
    let Some(t0) = t.first() else { return vec![] };
    let mut out: Vec<(String, usize)> = vec![];
    // two wires that carry their bus in the name of the wire, not in a prefix
    let bare = strip_digits(t0);
    if t.len() == 1 && (bare == "SCL" || bare == "SDA") { out.push((format!("I2C{}", &t0[bare.len()..]), 2)); }
    if t.len() == 1 && ["SWDIO", "SWCLK", "SWO", "SWDCLK"].contains(&t0.as_str()) { out.push(("SWD".into(), 2)); }
    if t.len() == 1 && ["TCK", "TMS", "TDI", "TDO", "NTRST"].contains(&t0.as_str()) { out.push(("JTAG".into(), 2)); }
    if t.len() >= 2 && t[1].chars().all(|c| c.is_ascii_digit()) { out.push((format!("{t0}-{}", t[1]), 2)); }
    // GHA, GHB, GHC: one letter tells the members apart
    if let Some((stem, _)) = lettered(name) { if family(&stem) { out.push((stem, 2)); } }
    if bare.len() != t0.len() { out.push((t0.clone(), 3)); }
    if bare.len() >= 2 { out.push((bare.to_string(), 2)); }
    out
}

/// Sort the nets into sections. `nets`: name and pins; `blocks`: blocks.json.
pub fn assign(nets: &[(String, Vec<String>)], blocks: &Value) -> Vec<Section> {
    let mut ground: Vec<String> = vec![]; let mut supply: Vec<(String, usize)> = vec![]; let mut rest: Vec<&(String, Vec<String>)> = vec![];
    for n in nets {
        if is_auto(&n.0) { rest.push(n); } else if is_ground(&n.0) { ground.push(n.0.clone()); } else if is_supply(&n.0) { supply.push((n.0.clone(), n.1.len())); } else { rest.push(n); }
    }
    let mut letters: HashMap<String, BTreeSet<char>> = HashMap::new();
    for n in rest.iter().filter(|n| !is_auto(&n.0)) { if let Some((stem, l)) = lettered(&n.0) { letters.entry(stem).or_default().insert(l); } }
    let family = |stem: &str| letters.get(stem).map(|l| l.len() >= 2 && FAMILIES.iter().any(|f| l.iter().all(|c| f.contains(*c)))).unwrap_or(false);
    let mut share: HashMap<String, usize> = HashMap::new();
    for n in rest.iter().filter(|n| !is_auto(&n.0)) { for (s, _) in stems(&n.0, &family) { *share.entry(s).or_default() += 1; } }
    let mut buses: BTreeMap<String, Vec<String>> = BTreeMap::new(); let mut left: Vec<&(String, Vec<String>)> = vec![];
    for n in rest {
        let hit = if is_auto(&n.0) { None } else { stems(&n.0, &family).into_iter().find(|(s, need)| share.get(s).copied().unwrap_or(0) >= *need) };
        match hit { Some((s, _)) => buses.entry(s).or_default().push(n.0.clone()), None => left.push(n) }
    }
    // what names do not place goes to the block that holds most of its pins
    let block_list: Vec<(String, Vec<String>)> = blocks["blocks"].as_array().into_iter().flatten().map(|b| (b["name"].as_str().or(b["id"].as_str()).unwrap_or("block").to_string(), b["parts"].as_array().into_iter().flatten().filter_map(|p| p.as_str().map(str::to_string)).collect())).collect();
    let block_of: HashMap<&str, usize> = block_list.iter().enumerate().flat_map(|(k, b)| b.1.iter().map(move |p| (p.as_str(), k))).collect();
    let mut by_block: Vec<Vec<String>> = vec![vec![]; block_list.len()]; let mut other: Vec<String> = vec![];
    for n in left {
        let mut votes = vec![0usize; block_list.len()];
        for p in &n.1 { if let Some(k) = block_of.get(p.split('.').next().unwrap_or("")) { votes[*k] += 1; } }
        match (0..votes.len()).filter(|k| votes[*k] > 0).max_by(|a, b| votes[*a].cmp(&votes[*b]).then(b.cmp(a))) { Some(k) => by_block[k].push(n.0.clone()), None => other.push(n.0.clone()) }
    }
    let sorted = |mut v: Vec<String>| { v.sort_by(|a, b| nat_cmp(a, b)); v };
    let mut out = vec![];
    if !ground.is_empty() { out.push(Section { kind: "ground", title: "Ground".into(), why: "named as a ground".into(), nets: sorted(ground) }); }
    supply.sort_by(|a, b| b.1.cmp(&a.1).then_with(|| nat_cmp(&a.0, &b.0)));
    for (n, _) in supply { out.push(Section { kind: "supply", title: format!("Supply rail {n}"), why: "named as a supply rail".into(), nets: vec![n] }); }
    let mut bus_list: Vec<(String, Vec<String>)> = buses.into_iter().collect();
    bus_list.sort_by(|a, b| nat_cmp(&a.0, &b.0));
    for (s, v) in bus_list { out.push(Section { kind: "bus", title: format!("{s} nets"), why: format!("names that share the stem {s}"), nets: sorted(v) }); }
    for (k, v) in by_block.into_iter().enumerate() { if !v.is_empty() { out.push(Section { kind: "block", title: format!("Other nets of the block: {}", block_list[k].0), why: "the name places it nowhere, and this block holds most of its pins".into(), nets: sorted(v) }); } }
    if !other.is_empty() { out.push(Section { kind: "other", title: "Other nets".into(), why: "no name pattern and no block".into(), nets: sorted(other) }); }
    out
}

#[cfg(test)]
mod tests {
    use super::*;
    use serde_json::json;
    fn n(name: &str, pins: &[&str]) -> (String, Vec<String>) { (name.to_string(), pins.iter().map(|p| p.to_string()).collect()) }
    fn find<'a>(s: &'a [Section], title: &str) -> &'a Section { s.iter().find(|x| x.title == title).unwrap_or_else(|| panic!("no section {title}: {:?}", s.iter().map(|x| &x.title).collect::<Vec<_>>())) }

    #[test]
    fn grounds_and_rails_are_told_from_signals_about_them() {
        for g in ["GND", "AGND", "GND_ISO", "VSS", "PWR_GND"] { assert!(is_ground(g), "{g}"); }
        for s in ["+3V3", "+3V3A", "+12V", "VIN", "VIN_LOGIC", "VBAT", "CM5_3.3V", "CM5_1.8V", "GPIO_VREF", "DRV_DVDD", "-5V"] { assert!(is_supply(s), "{s}"); }
        for s in ["VBUS_EN", "SD_VDD_OVERRIDE", "VSENSE", "SD_PWR_ON", "LDO5_ADJ", "BUCK_SW", "GPIO5", "HDMI1_TX2_P", "PHASE_X1", "WL_nDisable"] { assert!(!is_supply(s), "{s}"); assert!(!is_ground(s), "{s}"); }
        for a in ["N001", "CU_R4_1", "UNK_G3", "NET12"] { assert!(is_auto(a), "{a}"); }
        for a in ["NRST", "NTC", "NET_OK"] { assert!(!is_auto(a), "{a}"); }
    }

    #[test]
    fn buses_come_from_shared_stems_and_stay_apart_by_index() {
        let nets: Vec<(String, Vec<String>)> = ["GND", "+5V", "HDMI0_TX0_P", "HDMI0_TX0_N", "HDMI0_SDA", "HDMI1_TX0_P", "HDMI1_TX0_N", "HDMI1_SCL", "GPIO4", "GPIO17", "USB3-0-RX_N", "USB3-0-RX_P", "USB3-1-RX_N", "USB3-1-RX_P", "USB_P", "USB_N",
            "SCL0", "SDA0", "GHA", "GHB", "GHC", "SWDIO", "SWCLK", "PHASE_X1", "PHASE_X2", "LED1_K", "LED2_K", "NRST"].iter().map(|x| n(x, &[])).collect();
        let s = assign(&nets, &json!({}));
        assert_eq!(s[0].kind, "ground"); assert_eq!(s[1].title, "Supply rail +5V");
        assert_eq!(find(&s, "HDMI0 nets").nets.len(), 3); assert_eq!(find(&s, "HDMI1 nets").nets.len(), 3);
        assert_eq!(find(&s, "GPIO nets").nets, vec!["GPIO4", "GPIO17"]);
        assert_eq!(find(&s, "USB3-0 nets").nets.len(), 2); assert_eq!(find(&s, "USB3-1 nets").nets.len(), 2); assert_eq!(find(&s, "USB nets").nets, vec!["USB_N", "USB_P"]);
        assert_eq!(find(&s, "I2C0 nets").nets, vec!["SCL0", "SDA0"]); assert_eq!(find(&s, "GH nets").nets.len(), 3); assert_eq!(find(&s, "SWD nets").nets.len(), 2);
        assert_eq!(find(&s, "PHASE nets").nets.len(), 2); assert_eq!(find(&s, "LED nets").nets.len(), 2);
        assert_eq!(find(&s, "Other nets").nets, vec!["NRST"]);
        assert_eq!(s.iter().map(|x| x.nets.len()).sum::<usize>(), nets.len(), "every net is in one section");
        // a letter only makes a family when it is one of a set: NTC and NTC_ADC are the NTC nets, CAN and CAM stay apart
        let s = assign(&["NTC", "NTC_ADC", "CAN_H", "CAN_L", "CAM_GPIO0", "CAM_GPIO1", "SOA", "SOA_ADC", "SOB", "SOB_ADC"].iter().map(|x| n(x, &[])).collect::<Vec<_>>(), &json!({}));
        assert_eq!(find(&s, "NTC nets").nets.len(), 2); assert_eq!(find(&s, "CAN nets").nets.len(), 2); assert_eq!(find(&s, "CAM nets").nets.len(), 2); assert_eq!(find(&s, "SO nets").nets.len(), 4);
    }

    #[test]
    fn a_name_that_says_nothing_goes_to_the_block_holding_most_of_its_pins() {
        let blocks = json!({"blocks": [{"id": "mcu", "name": "Microcontroller", "parts": ["U1", "C1"]}, {"id": "can", "name": "CAN", "parts": ["U6", "J2"]}]});
        let s = assign(&[n("N001", &["U1.3", "C1.1", "U6.2"]), n("NRST", &["U6.1"]), n("N002", &["X9.1"])], &blocks);
        assert_eq!(find(&s, "Other nets of the block: Microcontroller").nets, vec!["N001"]);
        assert_eq!(find(&s, "Other nets of the block: CAN").nets, vec!["NRST"]);
        assert_eq!(find(&s, "Other nets").nets, vec!["N002"]);
    }
}