Optimize linkchecker by caching all filesystem access.

This should improve performance 2-3x depending on the system.

1 files changed, 214 insertions, 172 deletions

diff --git a/src/tools/linkchecker/main.rs b/src/tools/linkchecker/main.rs
index 7df4f5a9c46..076b3653583 100644
--- a/src/tools/linkchecker/main.rs
+++ b/src/tools/linkchecker/main.rs
@@ -14,10 +14,11 @@
 //! A few exceptions are allowed as there's known bugs in rustdoc, but this
 //! should catch the majority of "broken link" cases.
 
-use std::collections::hash_map::Entry;
+use std::cell::RefCell;
 use std::collections::{HashMap, HashSet};
 use std::env;
 use std::fs;
+use std::io::ErrorKind;
 use std::path::{Component, Path, PathBuf};
 use std::rc::Rc;
 use std::time::Instant;
@@ -25,8 +26,6 @@ use std::time::Instant;
 use once_cell::sync::Lazy;
 use regex::Regex;
 
-use crate::Redirect::*;
-
 // Add linkcheck exceptions here
 // If at all possible you should use intra-doc links to avoid linkcheck issues. These
 // are cases where that does not work
@@ -88,11 +87,10 @@ macro_rules! t {
 }
 
 fn main() {
-    let docs = env::args_os().nth(1).unwrap();
+    let docs = env::args_os().nth(1).expect("doc path should be first argument");
     let docs = env::current_dir().unwrap().join(docs);
-    let mut checker = Checker {
-        root: docs.clone(),
-        cache: HashMap::new(),
+    let mut checker = Checker { root: docs.clone(), cache: HashMap::new() };
+    let mut report = Report {
         errors: 0,
         start: Instant::now(),
         html_files: 0,
@@ -102,9 +100,9 @@ fn main() {
         links_ignored_exception: 0,
         intra_doc_exceptions: 0,
     };
-    checker.walk(&docs);
-    checker.report();
-    if checker.errors != 0 {
+    checker.walk(&docs, &mut report);
+    report.report();
+    if report.errors != 0 {
         println!("found some broken links");
         std::process::exit(1);
     }
@@ -113,6 +111,9 @@ fn main() {
 struct Checker {
     root: PathBuf,
     cache: Cache,
+}
+
+struct Report {
     errors: u32,
     start: Instant,
     html_files: u32,
@@ -123,23 +124,27 @@ struct Checker {
     intra_doc_exceptions: u32,
 }
 
-#[derive(Debug)]
-pub enum LoadError {
-    BrokenRedirect(PathBuf, std::io::Error),
-    IsRedirect,
+/// A cache entry.
+enum FileEntry {
+    /// An HTML file.
+    ///
+    /// This includes the contents of the HTML file, and an optional set of
+    /// HTML IDs. The IDs are used for checking fragments. The are computed
+    /// as-needed. The source is discarded (replaced with an empty string)
+    /// after the file has been checked, to conserve on memory.
+    HtmlFile { source: Rc<String>, ids: RefCell<HashSet<String>> },
+    /// This file is an HTML redirect to the given local path.
+    Redirect { target: PathBuf },
+    /// This is not an HTML file.
+    OtherFile,
+    /// This is a directory.
+    Dir,
+    /// The file doesn't exist.
+    Missing,
 }
 
-enum Redirect {
-    SkipRedirect,
-    FromRedirect(bool),
-}
-
-struct FileEntry {
-    source: Rc<String>,
-    ids: HashSet<String>,
-}
-
-type Cache = HashMap<PathBuf, FileEntry>;
+/// A cache to speed up file access.
+type Cache = HashMap<String, FileEntry>;
 
 fn small_url_encode(s: &str) -> String {
     s.replace("<", "%3C")
@@ -156,62 +161,36 @@ fn small_url_encode(s: &str) -> String {
         .replace("\"", "%22")
 }
 
-impl FileEntry {
-    fn parse_ids(&mut self, file: &Path, contents: &str, errors: &mut u32) {
-        if self.ids.is_empty() {
-            with_attrs_in_source(contents, " id", |fragment, i, _| {
-                let frag = fragment.trim_start_matches("#").to_owned();
-                let encoded = small_url_encode(&frag);
-                if !self.ids.insert(frag) {
-                    *errors += 1;
-                    println!("{}:{}: id is not unique: `{}`", file.display(), i, fragment);
-                }
-                // Just in case, we also add the encoded id.
-                self.ids.insert(encoded);
-            });
-        }
-    }
-}
-
 impl Checker {
-    fn walk(&mut self, dir: &Path) {
+    /// Primary entry point for walking the filesystem to find HTML files to check.
+    fn walk(&mut self, dir: &Path, report: &mut Report) {
         for entry in t!(dir.read_dir()).map(|e| t!(e)) {
             let path = entry.path();
             let kind = t!(entry.file_type());
             if kind.is_dir() {
-                self.walk(&path);
+                self.walk(&path, report);
             } else {
-                let pretty_path = self.check(&path);
-                if let Some(pretty_path) = pretty_path {
-                    let entry = self.cache.get_mut(&pretty_path).unwrap();
-                    // we don't need the source anymore,
-                    // so drop to reduce memory-usage
-                    entry.source = Rc::new(String::new());
-                }
+                self.check(&path, report);
             }
         }
     }
 
-    fn check(&mut self, file: &Path) -> Option<PathBuf> {
-        // Ignore non-HTML files.
-        if file.extension().and_then(|s| s.to_str()) != Some("html") {
-            return None;
-        }
-        self.html_files += 1;
-
-        let res = self.load_file(file, SkipRedirect);
-        let (pretty_file, contents) = match res {
-            Ok(res) => res,
-            Err(_) => return None,
+    /// Checks a single file.
+    fn check(&mut self, file: &Path, report: &mut Report) {
+        let (pretty_path, entry) = self.load_file(file, report);
+        let source = match entry {
+            FileEntry::Missing => panic!("missing file {:?} while walking", file),
+            FileEntry::Dir => unreachable!("never with `check` path"),
+            FileEntry::OtherFile => return,
+            FileEntry::Redirect { .. } => return,
+            FileEntry::HtmlFile { source, ids } => {
+                parse_ids(&mut ids.borrow_mut(), &pretty_path, source, report);
+                source.clone()
+            }
         };
-        self.cache.get_mut(&pretty_file).unwrap().parse_ids(
-            &pretty_file,
-            &contents,
-            &mut self.errors,
-        );
 
         // Search for anything that's the regex 'href[ ]*=[ ]*".*?"'
-        with_attrs_in_source(&contents, " href", |url, i, base| {
+        with_attrs_in_source(&source, " href", |url, i, base| {
             // Ignore external URLs
             if url.starts_with("http:")
                 || url.starts_with("https:")
@@ -220,10 +199,10 @@ impl Checker {
                 || url.starts_with("irc:")
                 || url.starts_with("data:")
             {
-                self.links_ignored_external += 1;
+                report.links_ignored_external += 1;
                 return;
             }
-            self.links_checked += 1;
+            report.links_checked += 1;
             let (url, fragment) = match url.split_once('#') {
                 None => (url, None),
                 Some((url, fragment)) => (url, Some(fragment)),
@@ -242,10 +221,10 @@ impl Checker {
                             // Avoid absolute paths as they make the docs not
                             // relocatable by making assumptions on where the docs
                             // are hosted relative to the site root.
-                            self.errors += 1;
+                            report.errors += 1;
                             println!(
                                 "{}:{}: absolute path - {}",
-                                pretty_file.display(),
+                                pretty_path,
                                 i + 1,
                                 Path::new(base).join(url).display()
                             );
@@ -262,138 +241,165 @@ impl Checker {
                 }
             }
 
-            // Alright, if we've found a file name then this file had better
-            // exist! If it doesn't then we register and print an error.
-            if path.exists() {
-                if path.is_dir() {
+            let (target_pretty_path, target_entry) = self.load_file(&path, report);
+            let (target_source, target_ids) = match target_entry {
+                FileEntry::Missing => {
+                    if is_exception(file, &target_pretty_path) {
+                        report.links_ignored_exception += 1;
+                    } else {
+                        report.errors += 1;
+                        println!(
+                            "{}:{}: broken link - `{}`",
+                            pretty_path,
+                            i + 1,
+                            target_pretty_path
+                        );
+                    }
+                    return;
+                }
+                FileEntry::Dir => {
                     // Links to directories show as directory listings when viewing
                     // the docs offline so it's best to avoid them.
-                    self.errors += 1;
-                    let pretty_path = path.strip_prefix(&self.root).unwrap_or(&path);
+                    report.errors += 1;
                     println!(
-                        "{}:{}: directory link - {}",
-                        pretty_file.display(),
+                        "{}:{}: directory link to `{}` \
+                         (directory links should use index.html instead)",
+                        pretty_path,
                         i + 1,
-                        pretty_path.display()
+                        target_pretty_path
                     );
                     return;
                 }
-                if let Some(extension) = path.extension() {
-                    // Ignore none HTML files.
-                    if extension != "html" {
-                        return;
+                FileEntry::OtherFile => return,
+                FileEntry::Redirect { target } => {
+                    let t = target.clone();
+                    drop(target);
+                    let (target, redir_entry) = self.load_file(&t, report);
+                    match redir_entry {
+                        FileEntry::Missing => {
+                            report.errors += 1;
+                            println!(
+                                "{}:{}: broken redirect from `{}` to `{}`",
+                                pretty_path,
+                                i + 1,
+                                target_pretty_path,
+                                target
+                            );
+                            return;
+                        }
+                        FileEntry::Redirect { target } => {
+                            // Redirect to a redirect, this link checker
+                            // currently doesn't support this, since it would
+                            // require cycle checking, etc.
+                            report.errors += 1;
+                            println!(
+                                "{}:{}: redirect from `{}` to `{}` \
+                                 which is also a redirect (not supported)",
+                                pretty_path,
+                                i + 1,
+                                target_pretty_path,
+                                target.display()
+                            );
+                            return;
+                        }
+                        FileEntry::Dir => {
+                            report.errors += 1;
+                            println!(
+                                "{}:{}: redirect from `{}` to `{}` \
+                                 which is a directory \
+                                 (directory links should use index.html instead)",
+                                pretty_path,
+                                i + 1,
+                                target_pretty_path,
+                                target
+                            );
+                            return;
+                        }
+                        FileEntry::OtherFile => return,
+                        FileEntry::HtmlFile { source, ids } => (source, ids),
                     }
                 }
-                let res = self.load_file(&path, FromRedirect(false));
-                let (pretty_path, contents) = match res {
-                    Ok(res) => res,
-                    Err(LoadError::BrokenRedirect(target, _)) => {
-                        self.errors += 1;
-                        println!(
-                            "{}:{}: broken redirect to {}",
-                            pretty_file.display(),
-                            i + 1,
-                            target.display()
-                        );
-                        return;
-                    }
-                    Err(LoadError::IsRedirect) => unreachable!(),
-                };
+                FileEntry::HtmlFile { source, ids } => (source, ids),
+            };
 
-                if let Some(ref fragment) = fragment {
-                    // Fragments like `#1-6` are most likely line numbers to be
-                    // interpreted by javascript, so we're ignoring these
-                    if fragment.splitn(2, '-').all(|f| f.chars().all(|c| c.is_numeric())) {
-                        return;
-                    }
+            // Alright, if we've found an HTML file for the target link. If
+            // this is a fragment link, also check that the `id` exists.
+            if let Some(ref fragment) = fragment {
+                // Fragments like `#1-6` are most likely line numbers to be
+                // interpreted by javascript, so we're ignoring these
+                if fragment.splitn(2, '-').all(|f| f.chars().all(|c| c.is_numeric())) {
+                    return;
+                }
 
-                    // These appear to be broken in mdbook right now?
-                    if fragment.starts_with('-') {
-                        return;
-                    }
+                // These appear to be broken in mdbook right now?
+                if fragment.starts_with('-') {
+                    return;
+                }
 
-                    let entry = self.cache.get_mut(&pretty_path).unwrap();
-                    entry.parse_ids(&pretty_path, &contents, &mut self.errors);
+                parse_ids(&mut target_ids.borrow_mut(), &pretty_path, target_source, report);
 
-                    if entry.ids.contains(*fragment) {
-                        return;
-                    }
-
-                    if is_exception(file, &format!("#{}", fragment)) {
-                        self.links_ignored_exception += 1;
-                    } else {
-                        self.errors += 1;
-                        print!("{}:{}: broken link fragment ", pretty_file.display(), i + 1);
-                        println!("`#{}` pointing to `{}`", fragment, pretty_path.display());
-                    };
+                if target_ids.borrow().contains(*fragment) {
+                    return;
                 }
-            } else {
-                let pretty_path = path.strip_prefix(&self.root).unwrap_or(&path);
-                if is_exception(file, pretty_path.to_str().unwrap()) {
+
+                if is_exception(file, &format!("#{}", fragment)) {
+                    report.links_ignored_exception += 1;
                 } else {
-                    self.errors += 1;
-                    print!("{}:{}: broken link - ", pretty_file.display(), i + 1);
-                    println!("{}", pretty_path.display());
-                }
+                    report.errors += 1;
+                    print!("{}:{}: broken link fragment ", pretty_path, i + 1);
+                    println!("`#{}` pointing to `{}`", fragment, pretty_path);
+                };
             }
         });
 
         // Search for intra-doc links that rustdoc didn't warn about
         // FIXME(#77199, 77200) Rustdoc should just warn about these directly.
         // NOTE: only looks at one line at a time; in practice this should find most links
-        for (i, line) in contents.lines().enumerate() {
+        for (i, line) in source.lines().enumerate() {
             for broken_link in BROKEN_INTRA_DOC_LINK.captures_iter(line) {
                 if is_intra_doc_exception(file, &broken_link[1]) {
-                    self.intra_doc_exceptions += 1;
+                    report.intra_doc_exceptions += 1;
                 } else {
-                    self.errors += 1;
-                    print!("{}:{}: broken intra-doc link - ", pretty_file.display(), i + 1);
+                    report.errors += 1;
+                    print!("{}:{}: broken intra-doc link - ", pretty_path, i + 1);
                     println!("{}", &broken_link[0]);
                 }
             }
         }
-        Some(pretty_file)
+        // we don't need the source anymore,
+        // so drop to reduce memory-usage
+        match self.cache.get_mut(&pretty_path).unwrap() {
+            FileEntry::HtmlFile { source, .. } => *source = Rc::new(String::new()),
+            _ => unreachable!("must be html file"),
+        }
     }
 
-    fn load_file(
-        &mut self,
-        file: &Path,
-        redirect: Redirect,
-    ) -> Result<(PathBuf, Rc<String>), LoadError> {
-        let pretty_file = PathBuf::from(file.strip_prefix(&self.root).unwrap_or(&file));
-
-        let (maybe_redirect, contents) = match self.cache.entry(pretty_file.clone()) {
-            Entry::Occupied(entry) => (None, entry.get().source.clone()),
-            Entry::Vacant(entry) => {
-                let contents = match fs::read_to_string(file) {
-                    Ok(s) => Rc::new(s),
-                    Err(err) => {
-                        return Err(if let FromRedirect(true) = redirect {
-                            LoadError::BrokenRedirect(file.to_path_buf(), err)
-                        } else {
-                            panic!("error loading {}: {}", file.display(), err);
-                        });
-                    }
-                };
-
-                let maybe = maybe_redirect(&contents);
-                if maybe.is_some() {
-                    self.html_redirects += 1;
-                    if let SkipRedirect = redirect {
-                        return Err(LoadError::IsRedirect);
+    /// Load a file from disk, or from the cache if available.
+    fn load_file(&mut self, file: &Path, report: &mut Report) -> (String, &FileEntry) {
+        let pretty_path =
+            file.strip_prefix(&self.root).unwrap_or(&file).to_str().unwrap().to_string();
+
+        let entry =
+            self.cache.entry(pretty_path.clone()).or_insert_with(|| match fs::metadata(file) {
+                Ok(metadata) if metadata.is_dir() => FileEntry::Dir,
+                Ok(_) => {
+                    if file.extension().and_then(|s| s.to_str()) != Some("html") {
+                        FileEntry::OtherFile
+                    } else {
+                        report.html_files += 1;
+                        load_html_file(file, report)
                     }
-                } else {
-                    entry.insert(FileEntry { source: contents.clone(), ids: HashSet::new() });
                 }
-                (maybe, contents)
-            }
-        };
-        match maybe_redirect.map(|url| file.parent().unwrap().join(url)) {
-            Some(redirect_file) => self.load_file(&redirect_file, FromRedirect(true)),
-            None => Ok((pretty_file, contents)),
-        }
+                Err(e) if e.kind() == ErrorKind::NotFound => FileEntry::Missing,
+                Err(e) => {
+                    panic!("unexpected read error for {}: {}", file.display(), e);
+                }
+            });
+        (pretty_path, entry)
     }
+}
 
+impl Report {
     fn report(&self) {
         println!("checked links in: {:.1}s", self.start.elapsed().as_secs_f64());
         println!("number of HTML files scanned: {}", self.html_files);
@@ -406,6 +412,25 @@ impl Checker {
     }
 }
 
+fn load_html_file(file: &Path, report: &mut Report) -> FileEntry {
+    let source = match fs::read_to_string(file) {
+        Ok(s) => Rc::new(s),
+        Err(err) => {
+            // This usually should not fail since `metadata` was already
+            // called successfully on this file.
+            panic!("unexpected read error for {}: {}", file.display(), err);
+        }
+    };
+    match maybe_redirect(&source) {
+        Some(target) => {
+            report.html_redirects += 1;
+            let target = file.parent().unwrap().join(target);
+            FileEntry::Redirect { target }
+        }
+        None => FileEntry::HtmlFile { source: source.clone(), ids: RefCell::new(HashSet::new()) },
+    }
+}
+
 fn is_intra_doc_exception(file: &Path, link: &str) -> bool {
     if let Some(entry) = INTRA_DOC_LINK_EXCEPTIONS.iter().find(|&(f, _)| file.ends_with(f)) {
         entry.1.is_empty() || entry.1.contains(&link)
@@ -432,6 +457,8 @@ fn is_exception(file: &Path, link: &str) -> bool {
     }
 }
 
+/// If the given HTML file contents is an HTML redirect, this returns the
+/// destination path given in the redirect.
 fn maybe_redirect(source: &str) -> Option<String> {
     const REDIRECT: &str = "<p>Redirecting to <a href=";
 
@@ -445,9 +472,9 @@ fn maybe_redirect(source: &str) -> Option<String> {
     })
 }
 
-fn with_attrs_in_source<F: FnMut(&str, usize, &str)>(contents: &str, attr: &str, mut f: F) {
+fn with_attrs_in_source<F: FnMut(&str, usize, &str)>(source: &str, attr: &str, mut f: F) {
     let mut base = "";
-    for (i, mut line) in contents.lines().enumerate() {
+    for (i, mut line) in source.lines().enumerate() {
         while let Some(j) = line.find(attr) {
             let rest = &line[j + attr.len()..];
             // The base tag should always be the first link in the document so
@@ -486,3 +513,18 @@ fn with_attrs_in_source<F: FnMut(&str, usize, &str)>(contents: &str, attr: &str,
         }
     }
 }
+
+fn parse_ids(ids: &mut HashSet<String>, file: &str, source: &str, report: &mut Report) {
+    if ids.is_empty() {
+        with_attrs_in_source(source, " id", |fragment, i, _| {
+            let frag = fragment.trim_start_matches("#").to_owned();
+            let encoded = small_url_encode(&frag);
+            if !ids.insert(frag) {
+                report.errors += 1;
+                println!("{}:{}: id is not unique: `{}`", file, i, fragment);
+            }
+            // Just in case, we also add the encoded id.
+            ids.insert(encoded);
+        });
+    }
+}