consortium_cli/
pinch.rs

1//! pinch — node set operations (replaces cluset/nodeset).
2//!
3//! Usage:
4//!   pinch -e node[1-5]          # expand: node1 node2 node3 node4 node5
5//!   pinch -c node[1-5]          # count: 5
6//!   pinch -f node1 node2 node3  # fold: node[1-3]
7//!   pinch -i node[1-5] node[3-7]  # intersection: node[3-5]
8//!   pinch -x node[3-4] node[1-5]  # difference: node[1-2,5]
9
10use std::io::{self, BufRead, Write};
11use std::process;
12
13use clap::Parser;
14
15use crate::output::{CliOutput, OutputArgs};
16use consortium::node_set::NodeSet;
17use consortium::range_set::RangeSet;
18
19/// pinch — node set operations for consortium.
20///
21/// Manipulate cluster node sets: fold, expand, count, and perform
22/// set algebra (union, intersection, difference, symmetric difference).
23#[derive(Parser)]
24#[command(name = "pinch", version, about)]
25pub struct Args {
26    // ── Operations (mutually exclusive) ────────────────────────────────
27    /// Count nodes in the resulting set.
28    #[arg(short = 'c', long = "count", group = "operation")]
29    count: bool,
30
31    /// Expand node set into individual node names.
32    #[arg(short = 'e', long = "expand", group = "operation")]
33    expand: bool,
34
35    /// Fold node names into compact bracket notation (default).
36    #[arg(short = 'f', long = "fold", group = "operation")]
37    fold: bool,
38
39    /// List node groups. Repeat for more detail (-l, -ll, -lll).
40    #[arg(short = 'l', long = "list", action = clap::ArgAction::Count, group = "operation")]
41    list: u8,
42
43    /// Regroup nodes using group definitions.
44    #[arg(short = 'r', long = "regroup", group = "operation")]
45    regroup: bool,
46
47    /// List available group sources.
48    #[arg(long = "groupsources", group = "operation")]
49    groupsources: bool,
50
51    // ── Set modifiers ──────────────────────────────────────────────────
52    /// Exclude nodes (set difference).
53    #[arg(short = 'x', long = "exclude")]
54    exclude: Vec<String>,
55
56    /// Intersect with node set.
57    #[arg(short = 'i', long = "intersection")]
58    intersection: Vec<String>,
59
60    /// Symmetric difference (XOR) with node set.
61    #[arg(short = 'X', long = "xor")]
62    xor: Vec<String>,
63
64    // ── Node sources ───────────────────────────────────────────────────
65    /// Use all nodes from default group.
66    #[arg(short = 'a', long = "all")]
67    all: bool,
68
69    /// Use nodes from named group.
70    #[arg(short = 'g', long = "group")]
71    group: Vec<String>,
72
73    // ── Display options ────────────────────────────────────────────────
74    /// Output separator (default: space for expand, newline for list).
75    #[arg(short = 'S', long = "separator")]
76    separator: Option<String>,
77
78    /// Operate on RangeSet instead of NodeSet.
79    #[arg(short = 'R', long = "rangeset")]
80    rangeset: bool,
81
82    /// Enable autostep folding (e.g., 2 → a-b/2).
83    #[arg(long = "autostep")]
84    autostep: Option<String>,
85
86    /// Split result into N subsets.
87    #[arg(long = "split")]
88    split: Option<usize>,
89
90    /// Split into contiguous subsets.
91    #[arg(long = "contiguous")]
92    contiguous: bool,
93
94    /// Pick N random nodes from result.
95    #[arg(long = "pick")]
96    pick: Option<usize>,
97
98    /// Pick node(s) at the given index position(s) (reverse of --index).
99    #[arg(short = 'I', long = "slice", value_name = "RANGESET")]
100    slice: Option<String>,
101
102    /// Output the index of NODE in the nodeset.
103    #[arg(long = "index", value_name = "NODE")]
104    index: Option<String>,
105
106    /// Output format string (%-style, applied to --index result).
107    #[arg(short = 'O', long = "output-format", value_name = "FORMAT")]
108    output_format: Option<String>,
109
110    // ── Positional ─────────────────────────────────────────────────────
111    /// Input node sets or node names. Reads from stdin if none given.
112    nodesets: Vec<String>,
113
114    #[command(flatten)]
115    output: OutputArgs,
116}
117
118pub fn run() {
119    let args = Args::parse();
120    let _out = CliOutput::from_args(&args.output);
121    if let Err(e) = execute(args) {
122        eprintln!("pinch: {e}");
123        process::exit(1);
124    }
125}
126
127fn execute(args: Args) -> anyhow::Result<()> {
128    // ── Command validation (mirrors CLI/Nodeset.py parser.error paths) ──
129    // --index counts as a command; combined with another command upstream
130    // reports "Multiple commands not allowed." and exits 2.
131    let cmdcount = args.count as u32
132        + args.expand as u32
133        + args.fold as u32
134        + u32::from(args.list > 0)
135        + args.regroup as u32
136        + args.groupsources as u32
137        + u32::from(args.index.is_some());
138    if cmdcount > 1 {
139        eprintln!("Multiple commands not allowed.");
140        process::exit(2);
141    }
142    if args.index.is_some() && args.pick.is_some() {
143        eprintln!("--index cannot be combined with --pick");
144        process::exit(2);
145    }
146
147    let stdout = io::stdout();
148    let mut out = stdout.lock();
149
150    if args.rangeset {
151        return run_rangeset(&args, &mut out);
152    }
153
154    // ── Build the base node set ────────────────────────────────────────
155    let mut ns = gather_input(&args)?;
156
157    // ── Apply set modifiers ────────────────────────────────────────────
158    for pat in &args.exclude {
159        let other = NodeSet::parse(pat)?;
160        ns.difference_update(&other);
161    }
162    for pat in &args.intersection {
163        let other = NodeSet::parse(pat)?;
164        ns.intersection_update(&other);
165    }
166    for pat in &args.xor {
167        let other = NodeSet::parse(pat)?;
168        ns.symmetric_difference_update(&other);
169    }
170
171    // ── Slice transform (-I), applied before --pick and --index ────────
172    if let Some(ref slice_pat) = args.slice {
173        ns = slice_nodeset(&ns, slice_pat)?;
174    }
175
176    // ── Pick random subset ─────────────────────────────────────────────
177    if let Some(n) = args.pick {
178        ns = pick_random(&ns, n);
179    }
180
181    let fmt = args.output_format.as_deref().unwrap_or("%s");
182
183    // ── --index: position of a node in the set (reverse of -I/--slice) ──
184    if let Some(ref node_arg) = args.index {
185        // a single node is required (like list.index())
186        match NodeSet::parse(node_arg) {
187            Ok(parsed) if parsed.len() == 1 => {}
188            Ok(_) => {
189                eprintln!("ERROR: index() argument must be a single node");
190                process::exit(1);
191            }
192            Err(e) => return Err(e.into()),
193        }
194        return match ns.index(node_arg) {
195            Some(i) => {
196                writeln!(out, "{}", apply_format(fmt, &i.to_string()))?;
197                Ok(())
198            }
199            None => {
200                eprintln!("ERROR: '{node_arg}' is not in nodeset");
201                process::exit(1);
202            }
203        };
204    }
205
206    // ── Output ─────────────────────────────────────────────────────────
207    let sep = args.separator.as_deref().unwrap_or(" ");
208
209    if args.count {
210        writeln!(out, "{}", ns.len())?;
211    } else if args.expand {
212        let nodes: Vec<String> = ns.iter().collect();
213        writeln!(out, "{}", nodes.join(sep))?;
214    } else if let Some(n) = args.split {
215        let subsets = ns.split(n);
216        for subset in subsets {
217            writeln!(out, "{subset}")?;
218        }
219    } else if args.contiguous {
220        // Split into contiguous ranges — each pattern separately
221        let subsets = ns.split(ns.len().max(1));
222        for subset in subsets {
223            if !subset.is_empty() {
224                writeln!(out, "{subset}")?;
225            }
226        }
227    } else {
228        // Default: fold
229        writeln!(out, "{ns}")?;
230    }
231
232    Ok(())
233}
234
235/// Collect input nodesets from args + stdin.
236fn gather_input(args: &Args) -> anyhow::Result<NodeSet> {
237    let mut ns = NodeSet::new();
238
239    // From positional args
240    for pat in &args.nodesets {
241        let parsed = NodeSet::parse(pat)?;
242        ns.update(&parsed);
243    }
244
245    // From stdin if no positional args and not a tty
246    if args.nodesets.is_empty() && !is_terminal::is_terminal(io::stdin()) {
247        let stdin = io::stdin();
248        for line in stdin.lock().lines() {
249            let line = line?;
250            let trimmed = line.trim();
251            if !trimmed.is_empty() {
252                let parsed = NodeSet::parse(trimmed)?;
253                ns.update(&parsed);
254            }
255        }
256    }
257
258    Ok(ns)
259}
260
261/// RangeSet mode — operates on numeric ranges, not node names.
262fn run_rangeset(args: &Args, out: &mut impl Write) -> anyhow::Result<()> {
263    let autostep = parse_autostep(&args.autostep);
264
265    let mut rs = RangeSet::new();
266
267    // From positional args
268    for pat in &args.nodesets {
269        let parsed = RangeSet::parse(pat, autostep)?;
270        rs.update(&parsed);
271    }
272
273    // From stdin
274    if args.nodesets.is_empty() && !is_terminal::is_terminal(io::stdin()) {
275        let stdin = io::stdin();
276        for line in stdin.lock().lines() {
277            let line = line?;
278            let trimmed = line.trim();
279            if !trimmed.is_empty() {
280                let parsed = RangeSet::parse(trimmed, autostep)?;
281                rs.update(&parsed);
282            }
283        }
284    }
285
286    // Apply set modifiers
287    for pat in &args.exclude {
288        let other = RangeSet::parse(pat, autostep)?;
289        rs.difference_update(&other);
290    }
291    for pat in &args.intersection {
292        let other = RangeSet::parse(pat, autostep)?;
293        rs.intersection_update(&other);
294    }
295    for pat in &args.xor {
296        let other = RangeSet::parse(pat, autostep)?;
297        rs.symmetric_difference_update(&other);
298    }
299
300    // Slice transform (-I), applied before --index like upstream
301    if let Some(ref slice_pat) = args.slice {
302        rs = slice_rangeset(&rs, slice_pat)?;
303    }
304
305    let fmt = args.output_format.as_deref().unwrap_or("%s");
306
307    // --index: position of an element in the rangeset (padding significant)
308    if let Some(ref elem_arg) = args.index {
309        return match rs.index_str(elem_arg) {
310            Some(i) => {
311                writeln!(out, "{}", apply_format(fmt, &i.to_string()))?;
312                Ok(())
313            }
314            None => {
315                eprintln!("ERROR: {elem_arg} is not in RangeSet");
316                process::exit(1);
317            }
318        };
319    }
320
321    let sep = args.separator.as_deref().unwrap_or(" ");
322
323    if args.count {
324        writeln!(out, "{}", rs.len())?;
325    } else if args.expand {
326        let items: Vec<String> = rs.striter().collect();
327        writeln!(out, "{}", items.join(sep))?;
328    } else {
329        // fold
330        writeln!(out, "{rs}")?;
331    }
332
333    Ok(())
334}
335
336/// Parse autostep argument: a number, "auto", or "%N".
337fn parse_autostep(s: &Option<String>) -> Option<u32> {
338    let s = s.as_deref()?;
339    match s {
340        "auto" => Some(0), // 0 means auto-detect
341        _ => s.parse::<u32>().ok(),
342    }
343}
344
345/// Apply the `-I/--slice` transform: keep the elements at the positions
346/// selected by the slice rangeset (upstream `xset[sli]` for each slice of
347/// `RangeSet(options.slice_rangeset).slices()`; out-of-range positions are
348/// clamped away like Python slicing).
349fn slice_nodeset(ns: &NodeSet, slice_pat: &str) -> anyhow::Result<NodeSet> {
350    let positions = RangeSet::parse(slice_pat, None)?;
351    let elems: Vec<String> = ns.iter().collect();
352    let mut sliced = NodeSet::new();
353    for pos in positions.intiter() {
354        if pos >= 0 && (pos as usize) < elems.len() {
355            sliced.update_str(&elems[pos as usize])?;
356        }
357    }
358    Ok(sliced)
359}
360
361/// RangeSet mode equivalent of [`slice_nodeset`].
362fn slice_rangeset(rs: &RangeSet, slice_pat: &str) -> anyhow::Result<RangeSet> {
363    let positions = RangeSet::parse(slice_pat, None)?;
364    let elems = rs.sorted();
365    let mut sliced = RangeSet::new();
366    for pos in positions.intiter() {
367        if pos >= 0 && (pos as usize) < elems.len() {
368            sliced.add_str(&elems[pos as usize]);
369        }
370    }
371    Ok(sliced)
372}
373
374/// Apply a Python `%`-style output format string to a single value
375/// (`%s`/`%d`/`%i` substitute the value, `%%` is a literal percent).
376fn apply_format(fmt: &str, value: &str) -> String {
377    let mut out = String::new();
378    let mut chars = fmt.chars();
379    while let Some(c) = chars.next() {
380        if c == '%' {
381            match chars.next() {
382                Some('s') | Some('d') | Some('i') => out.push_str(value),
383                Some('%') => out.push('%'),
384                Some(other) => {
385                    out.push('%');
386                    out.push(other);
387                }
388                None => out.push('%'),
389            }
390        } else {
391            out.push(c);
392        }
393    }
394    out
395}
396
397/// Pick N random nodes from a NodeSet.
398fn pick_random(ns: &NodeSet, n: usize) -> NodeSet {
399    use std::collections::HashSet;
400
401    let all: Vec<String> = ns.iter().collect();
402    if n >= all.len() {
403        return ns.clone();
404    }
405
406    // Simple Fisher-Yates-ish random pick using system entropy
407    let mut indices: Vec<usize> = (0..all.len()).collect();
408    // Use a simple LCG seeded from process id + time for shuffling
409    let mut seed: u64 = std::process::id() as u64
410        ^ std::time::SystemTime::now()
411            .duration_since(std::time::UNIX_EPOCH)
412            .unwrap_or_default()
413            .as_nanos() as u64;
414
415    for i in (1..indices.len()).rev() {
416        seed = seed.wrapping_mul(6364136223846793005).wrapping_add(1);
417        let j = (seed >> 33) as usize % (i + 1);
418        indices.swap(i, j);
419    }
420
421    let picked: HashSet<&str> = indices[..n].iter().map(|&i| all[i].as_str()).collect();
422    let mut result = NodeSet::new();
423    for node in picked {
424        let _ = result.update_str(node);
425    }
426    result
427}
consortium_cli/pinch.rs

consortium_cli/
pinch.rs
