master
Colby Knox Merge PR #351: kicad-lane: 1.0.88, kicad_variants and kicad_sch_zoom b960256 11d ago
//! KiCad 10's design variants (the standard's `variants`): which parts are fitted on which build.
//! The project names them (`.kicad_pro` `schematic.variants: [{name, description?}]`), and each placed symbol says how it
//! differs in a variant inside its instance path, writing only what differs from the symbol itself:
//! `(instances (project "p" (path "/…" (reference "R5") (unit 1) (variant (name "Lite") (dnp yes) (field (name "Value") (value "22k"))))))`.
//! A part is named by the reference its path gives it, so a sheet used twice has a reference (and a variation) per use.
use crate::schematic::{pretty, Node, Schematic};
use serde_json::{json, Value};
use std::collections::BTreeMap;

/// One part's variation in one variant.
#[derive(Clone, Debug, Default, PartialEq)]
pub struct Variation {
    pub reference: String,
    pub dnp: Option<bool>,
    pub in_bom: Option<bool>,
    pub on_board: Option<bool>,
    pub in_pos_files: Option<bool>,
    pub exclude_from_sim: Option<bool>,
    pub fields: Vec<(String, String)>,
}

fn flag(n: &Node, k: &str) -> Option<bool> {
    n.child_value(k).map(|v| v == "yes")
}

fn is_list(n: &Node, head: &str) -> bool {
    n.is_list() && n.head() == Some(head)
}

/// Every symbol instance path on the sheet: (top-level item index, path node index path) visited mutably.
fn for_each_path(sch: &mut Schematic, f: &mut dyn FnMut(&mut Node)) {
    for sym in sch.top_mut().iter_mut().filter(|n| crate::schematic::is_symbol_instance(n)) {
        let Some(inst) = sym.child_mut("instances") else { continue };
        for proj in inst.children_mut().into_iter().flatten().filter(|n| is_list(n, "project")) {
            for path in proj.children_mut().into_iter().flatten().filter(|n| is_list(n, "path")) {
                f(path);
            }
        }
    }
}

/// Every variation on the sheet, by variant name.
pub fn read(sch: &Schematic) -> BTreeMap<String, Vec<Variation>> {
    let mut out: BTreeMap<String, Vec<Variation>> = BTreeMap::new();
    for sym in sch.top().iter().filter(|n| crate::schematic::is_symbol_instance(n)) {
        let Some(inst) = sym.child("instances") else { continue };
        for proj in inst.children().iter().filter(|n| is_list(n, "project")) {
            for path in proj.children().iter().filter(|n| is_list(n, "path")) {
                let reference = path.child_value("reference").unwrap_or_default();
                for v in path.children().iter().filter(|n| is_list(n, "variant")) {
                    let Some(name) = v.child_value("name") else { continue };
                    let fields = v.children().iter().filter(|n| is_list(n, "field"))
                        .filter_map(|f| Some((f.child_value("name")?, f.child_value("value").unwrap_or_default()))).collect();
                    out.entry(name).or_default().push(Variation {
                        reference: reference.clone(), dnp: flag(v, "dnp"), in_bom: flag(v, "in_bom"), on_board: flag(v, "on_board"),
                        in_pos_files: flag(v, "in_pos_files"), exclude_from_sim: flag(v, "exclude_from_sim"), fields,
                    });
                }
            }
        }
    }
    out
}

/// Mark `reference` not fitted (DNP) in `variant`, adding the variation when the part has none there. Returns how many
/// instance paths changed (0 when the part is not on this sheet).
pub fn set_not_fitted(sch: &mut Schematic, reference: &str, variant: &str) -> usize {
    let mut n = 0;
    for_each_path(sch, &mut |path| {
        if path.child_value("reference").as_deref() != Some(reference) {
            return;
        }
        let Some(children) = path.children_mut() else { return };
        match children.iter_mut().find(|c| is_list(c, "variant") && c.child_value("name").as_deref() == Some(variant)) {
            Some(v) => {
                if let Some(d) = v.child_mut("dnp") {
                    d.set_arg(0, Node::atom("yes"));
                } else {
                    let mut d = Node::list("dnp", vec![Node::atom("yes")]);
                    pretty(&mut d, 6);
                    let at = v.children().iter().position(|c| is_list(c, "name")).map(|i| i + 1).unwrap_or(1);
                    if let Some(vc) = v.children_mut() {
                        vc.insert(at, d);
                    }
                }
            }
            None => {
                let mut v = Node::list("variant", vec![Node::list("name", vec![Node::quoted(variant)]), Node::list("dnp", vec![Node::atom("yes")])]);
                pretty(&mut v, 5);
                children.push(v);
            }
        }
        n += 1;
    });
    n
}

/// Put `reference` back in `variant`: its whole variation there is dropped (fitted as the symbol is). Returns how many
/// instance paths changed.
pub fn set_fitted(sch: &mut Schematic, reference: &str, variant: &str) -> usize {
    let mut n = 0;
    for_each_path(sch, &mut |path| {
        if path.child_value("reference").as_deref() != Some(reference) {
            return;
        }
        if let Some(children) = path.children_mut() {
            let before = children.len();
            children.retain(|c| !(is_list(c, "variant") && c.child_value("name").as_deref() == Some(variant)));
            n += before - children.len();
        }
    });
    n
}

/// Drop every part's variation in `variant`. Returns how many were dropped.
pub fn remove(sch: &mut Schematic, variant: &str) -> usize {
    let mut n = 0;
    for_each_path(sch, &mut |path| {
        if let Some(children) = path.children_mut() {
            let before = children.len();
            children.retain(|c| !(is_list(c, "variant") && c.child_value("name").as_deref() == Some(variant)));
            n += before - children.len();
        }
    });
    n
}

/// Every reference placed on the sheet (one per instance path).
pub fn references(sch: &Schematic) -> Vec<String> {
    let mut out = Vec::new();
    for sym in sch.top().iter().filter(|n| crate::schematic::is_symbol_instance(n)) {
        for proj in sym.child("instances").map(|i| i.children()).unwrap_or(&[]).iter().filter(|n| is_list(n, "project")) {
            out.extend(proj.children().iter().filter(|n| is_list(n, "path")).filter_map(|p| p.child_value("reference")));
        }
    }
    out
}

/// The project's variant list (`schematic.variants`).
pub fn project_variants(pro: &Value) -> Vec<(String, String)> {
    pro.pointer("/schematic/variants").and_then(Value::as_array).into_iter().flatten()
        .filter_map(|e| Some((e["name"].as_str()?.to_string(), e["description"].as_str().unwrap_or("").to_string()))).collect()
}

/// Add (or describe) a variant in the project; returns true when the project changed.
pub fn project_add(pro: &mut Value, name: &str, description: Option<&str>) -> bool {
    if !pro["schematic"].is_object() {
        pro["schematic"] = json!({});
    }
    if !pro["schematic"]["variants"].is_array() {
        pro["schematic"]["variants"] = json!([]);
    }
    let list = pro["schematic"]["variants"].as_array_mut().expect("array");
    match list.iter_mut().find(|e| e["name"].as_str() == Some(name)) {
        Some(e) => match description {
            Some(d) if e["description"].as_str().unwrap_or("") != d => {
                if d.is_empty() { e.as_object_mut().map(|o| o.remove("description")); } else { e["description"] = json!(d); }
                true
            }
            _ => false,
        },
        None => {
            let mut e = json!({"name": name});
            if let Some(d) = description.filter(|d| !d.is_empty()) {
                e["description"] = json!(d);
            }
            list.push(e);
            true
        }
    }
}

/// Remove a variant from the project; returns true when it was there.
pub fn project_remove(pro: &mut Value, name: &str) -> bool {
    let Some(list) = pro.pointer_mut("/schematic/variants").and_then(Value::as_array_mut) else { return false };
    let before = list.len();
    list.retain(|e| e["name"].as_str() != Some(name));
    before != list.len()
}

/// The standard's answer for every variant: the project's list joined with the sheets' variations.
pub fn summary(pro: &Value, sheets: &[&Schematic]) -> Vec<Value> {
    let mut by: BTreeMap<String, Vec<Variation>> = BTreeMap::new();
    for s in sheets {
        for (k, v) in read(s) {
            by.entry(k).or_default().extend(v);
        }
    }
    let mut names: Vec<(String, String)> = project_variants(pro);
    for k in by.keys() {
        if !names.iter().any(|(n, _)| n == k) {
            names.push((k.clone(), String::new()));
        }
    }
    names.into_iter().map(|(name, description)| {
        let vs = by.remove(&name).unwrap_or_default();
        let mut not_fitted: Vec<&str> = vs.iter().filter(|v| v.dnp == Some(true)).map(|v| v.reference.as_str()).collect();
        not_fitted.sort();
        not_fitted.dedup();
        let alternates: Vec<Value> = vs.iter().filter(|v| !v.fields.is_empty())
            .map(|v| json!({"designator": v.reference, "fields": v.fields.iter().map(|(k, x)| (k.clone(), json!(x))).collect::<serde_json::Map<_, _>>()})).collect();
        let other: Vec<Value> = vs.iter().filter(|v| v.in_bom.is_some() || v.on_board.is_some() || v.in_pos_files.is_some() || v.exclude_from_sim.is_some())
            .map(|v| json!({"designator": v.reference, "inBom": v.in_bom, "onBoard": v.on_board, "inPosFiles": v.in_pos_files, "excludeFromSim": v.exclude_from_sim})).collect();
        json!({"name": name, "description": description, "notFitted": not_fitted, "alternates": alternates, "flags": other,
               "inProject": project_variants(pro).iter().any(|(n, _)| *n == name)})
    }).collect()
}

/// The parts KiCad's own XML netlist marks not fitted in `variant`: DNP on the symbol itself, or a
/// `<variants><variant name="…"><property name="dnp" value="1"/>` block in the part (kicad-cli 10 writes variant
/// differences that way).
pub fn netlist_not_fitted(xml: &str, variant: &str) -> Vec<String> {
    let mut out = Vec::new();
    let mut rest = xml;
    while let Some(i) = rest.find("<comp ref=\"") {
        let after = &rest[i + 11..];
        let Some(q) = after.find('"') else { break };
        let reference = &after[..q];
        let end = after.find("</comp>").unwrap_or(after.len());
        let body = &after[..end];
        let (own, vars) = match body.find("<variants>") {
            Some(k) => (&body[..k], &body[k..body.find("</variants>").unwrap_or(body.len())]),
            None => (body, ""),
        };
        let base = own.contains("<property name=\"dnp\"/>") || own.contains("<property name=\"dnp\" value=\"1\"/>");
        let tag = format!("<variant name=\"{}\">", variant.replace('&', "&amp;").replace('"', "&quot;").replace('<', "&lt;"));
        let in_variant = vars.find(&tag).map(|k| {
            let v = &vars[k..];
            let v = &v[..v.find("</variant>").unwrap_or(v.len())];
            if v.contains("<property name=\"dnp\" value=\"1\"/>") { Some(true) } else if v.contains("<property name=\"dnp\" value=\"0\"/>") { Some(false) } else { None }
        }).flatten();
        if in_variant.unwrap_or(base) {
            out.push(reference.to_string());
        }
        rest = &after[end..];
    }
    out.sort();
    out.dedup();
    out
}

#[cfg(test)]
mod tests {
    #[test]
    fn kicads_netlist_says_which_parts_a_variant_leaves_out() {
        // as kicad-cli 10.0.7 writes it for the pic_programmer demo with C1 not fitted in Lite
        let xml = "<export><components><comp ref=\"C1\">\n<value>100uF</value>\n<property name=\"Sheetname\" value=\"p\"/>\n<variants>\n<variant name=\"Lite\">\n<property name=\"dnp\" value=\"1\"/>\n</variant>\n</variants>\n<property name=\"ki_fp_filters\" value=\"C*\"/>\n</comp>\n<comp ref=\"C2\">\n<value>1n</value>\n</comp>\n<comp ref=\"J9\">\n<property name=\"dnp\"/>\n</comp></components></export>";
        assert_eq!(super::netlist_not_fitted(xml, "Lite"), vec!["C1".to_string(), "J9".to_string()]);
        assert_eq!(super::netlist_not_fitted(xml, "Pro"), vec!["J9".to_string()]);
    }

    use super::*;

    const SCH: &str = "(kicad_sch\n\t(version 20260306)\n\t(generator \"eeschema\")\n\t(uuid \"aaaa\")\n\t(symbol\n\t\t(lib_id \"Device:R\")\n\t\t(at 10 10 0)\n\t\t(uuid \"r1\")\n\t\t(property \"Reference\" \"R1\"\n\t\t\t(at 0 0 0)\n\t\t)\n\t\t(instances\n\t\t\t(project \"p\"\n\t\t\t\t(path \"/aaaa\"\n\t\t\t\t\t(reference \"R1\")\n\t\t\t\t\t(unit 1)\n\t\t\t\t\t(variant\n\t\t\t\t\t\t(name \"Pro\")\n\t\t\t\t\t\t(field\n\t\t\t\t\t\t\t(name \"Value\")\n\t\t\t\t\t\t\t(value \"22k\")\n\t\t\t\t\t\t)\n\t\t\t\t\t)\n\t\t\t\t)\n\t\t\t)\n\t\t)\n\t)\n\t(symbol\n\t\t(lib_id \"Device:C\")\n\t\t(at 20 10 0)\n\t\t(uuid \"c1\")\n\t\t(property \"Reference\" \"C1\"\n\t\t\t(at 0 0 0)\n\t\t)\n\t\t(instances\n\t\t\t(project \"p\"\n\t\t\t\t(path \"/aaaa\"\n\t\t\t\t\t(reference \"C1\")\n\t\t\t\t\t(unit 1)\n\t\t\t\t)\n\t\t\t)\n\t\t)\n\t)\n)\n";

    #[test]
    fn variations_read_mark_put_back_and_remove() {
        let mut s = Schematic::parse(SCH).unwrap();
        let r = read(&s);
        assert_eq!(r["Pro"][0].fields, vec![("Value".to_string(), "22k".to_string())]);
        // C1 not fitted in Lite: a new variation, laid out the way eeschema writes it
        assert_eq!(set_not_fitted(&mut s, "C1", "Lite"), 1);
        let t = s.to_text();
        assert!(t.contains("\t\t\t\t\t(variant\n\t\t\t\t\t\t(name \"Lite\")\n\t\t\t\t\t\t(dnp yes)\n\t\t\t\t\t)"), "{t}");
        // R1 not fitted in Pro too: the existing variation gains dnp, keeping its field
        assert_eq!(set_not_fitted(&mut s, "R1", "Pro"), 1);
        let r = read(&Schematic::parse(&s.to_text()).unwrap());
        assert_eq!((r["Pro"][0].dnp, r["Pro"][0].fields.len()), (Some(true), 1));
        assert_eq!(r["Lite"][0].reference, "C1");
        assert_eq!(set_not_fitted(&mut s, "R9", "Lite"), 0);
        // putting C1 back drops its Lite variation
        assert_eq!(set_fitted(&mut s, "C1", "Lite"), 1);
        assert!(!read(&s).contains_key("Lite"));
        assert_eq!(remove(&mut s, "Pro"), 1);
        assert!(read(&s).is_empty());
        assert_eq!(references(&s), vec!["R1".to_string(), "C1".to_string()]);
    }

    #[test]
    fn the_project_lists_variants_by_name() {
        let mut pro = json!({"schematic": {"variants": []}, "meta": {}});
        assert!(project_add(&mut pro, "Lite", Some("no USB")));
        assert!(!project_add(&mut pro, "Lite", Some("no USB")));
        assert!(project_add(&mut pro, "Pro", None));
        assert_eq!(project_variants(&pro), vec![("Lite".into(), "no USB".into()), ("Pro".into(), String::new())]);
        let s = Schematic::parse(SCH).unwrap();
        let sum = summary(&pro, &[&s]);
        assert_eq!(sum[1]["alternates"], json!([{"designator": "R1", "fields": {"Value": "22k"}}]));
        assert!(project_remove(&mut pro, "Lite") && !project_remove(&mut pro, "Lite"));
    }
}