diff --git a/src/book.rs b/src/book.rs index 4a9ec98..b834ef2 100644 --- a/src/book.rs +++ b/src/book.rs @@ -310,6 +310,12 @@ pub fn load_epub(path: impl AsRef) -> Result { .cloned() .collect(); + let spine_paths: Vec = spine.iter().map(|id| { + doc.resources.get(id) + .map(|(path, _)| path.to_string_lossy().to_string()) + .unwrap_or_else(|| id.clone()) + }).collect(); + let mut sections = Vec::new(); for (i, href) in spine.iter().enumerate() { let raw_html = doc.get_resource_str(href) @@ -326,7 +332,7 @@ pub fn load_epub(path: impl AsRef) -> Result { } let raw_toc = std::mem::take(&mut doc.toc); - let toc = build_toc(&raw_toc, &spine); + let toc = build_toc(&raw_toc, &spine, &spine_paths); Ok(Book { title, author, cover, layout, sections, toc }) } @@ -351,8 +357,8 @@ fn extract_title(html: &str) -> Option { } fn extract_filename(path: &str) -> &str { - let path = path.trim_end_matches('/'); - path.rsplit('/').next().unwrap_or(path) + let path = path.trim_end_matches('/').trim_end_matches('\\'); + path.rsplit(&['/', '\\'][..]).next().unwrap_or(path) } fn extract_fragment(path: &str) -> Option { @@ -371,6 +377,7 @@ fn extract_fragment(path: &str) -> Option { fn build_toc( entries: &[epub::doc::NavPoint], spine: &[String], + spine_paths: &[String], ) -> Vec { entries .iter() @@ -378,20 +385,29 @@ fn build_toc( let content_str = e.content.to_string_lossy(); let anchor = extract_fragment(&content_str); let content_file = extract_filename(&content_str); - let section = spine + let section = spine_paths .iter() .position(|s| { let spine_file = extract_filename(s); - spine_file == content_file - || content_str.contains(s.as_str()) - || s.contains(content_str.as_ref()) + if spine_file == content_file { + return true; + } + content_str.contains(s.as_str()) || s.contains(content_str.as_ref()) + }) + .or_else(|| { + spine.iter().position(|s| { + let spine_file = extract_filename(s); + spine_file == content_file + || content_str.contains(s.as_str()) + || s.contains(content_str.as_ref()) + }) }) .unwrap_or(0); TocEntry { label: e.label.clone(), section, anchor, - children: build_toc(&e.children, spine), + children: build_toc(&e.children, spine, spine_paths), } }) .collect() @@ -521,7 +537,7 @@ mod tests { #[test] fn test_build_toc_empty() { - let toc = build_toc(&[], &[]); + let toc = build_toc(&[], &[], &[]); assert!(toc.is_empty()); } @@ -561,7 +577,7 @@ mod tests { children: vec![], }, ]; - let toc = build_toc(&nav_points, &spine); + let toc = build_toc(&nav_points, &spine, &spine); assert_eq!(toc.len(), 1); assert_eq!(toc[0].section, 1); // maps to spine index 1 assert_eq!(toc[0].label, "Chapter 2"); @@ -582,7 +598,7 @@ mod tests { children: vec![], }, ]; - let toc = build_toc(&nav_points, &spine); + let toc = build_toc(&nav_points, &spine, &spine); assert_eq!(toc.len(), 1); assert_eq!(toc[0].section, 0); }