rust-kernel / user / wc.rs

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
//! wc [-l] [file...] -- count lines, words and bytes (stdin if no files).
//! -l: only lines.
#![no_std]
#![no_main]

#[macro_use]
#[path = "ulib.rs"]
mod ulib;

use ulib::*;

/// (lines, words, bytes) of everything readable from `fd`.
fn count(fd: i64) -> Result<[u64; 3], i64> {
    let mut c = [0u64; 3];
    let mut in_word = false;
    let mut buf = [0u8; 512];
    loop {
        let n = read(fd, &mut buf);
        if n < 0 {
            return Err(n);
        }
        if n == 0 {
            return Ok(c);
        }
        for &b in &buf[..n as usize] {
            c[2] += 1;
            if b == b'\n' {
                c[0] += 1;
            }
            let space = b == b' ' || b == b'\n' || b == b'\t' || b == b'\r';
            if !space && !in_word {
                c[1] += 1; // a word starts here
            }
            in_word = !space;
        }
    }
}

fn show(c: [u64; 3], lines_only: bool, name: &str) {
    let sep = if name.is_empty() { "" } else { " " };
    if lines_only {
        println!("{}{}{}", c[0], sep, name);
    } else {
        println!("{:>7} {:>7} {:>7}{}{}", c[0], c[1], c[2], sep, name);
    }
}

fn main() -> i64 {
    let lines_only = arg(1) == Some("-l");
    let first = if lines_only { 2 } else { 1 };
    if argc() <= first {
        return match count(0) {
            Ok(c) => {
                show(c, lines_only, "");
                0
            }
            Err(_) => 1,
        };
    }
    let (mut total, mut status) = ([0u64; 3], 0);
    for i in first..argc() {
        let name = arg(i).unwrap_or("");
        let fd = open(name, 0);
        let r = if fd < 0 { Err(fd) } else { count(fd) };
        if fd >= 0 {
            close(fd);
        }
        match r {
            Ok(c) => {
                show(c, lines_only, name);
                (0..3).for_each(|k| total[k] += c[k]);
            }
            Err(e) => {
                eprintln!("wc: {}: {}", name, strerror(e));
                status = 1;
            }
        }
    }
    if argc() - first > 1 {
        show(total, lines_only, "total");
    }
    status
}