Load themes lazily for improved startup time (#1969)

This is for #951. Syntax lazy-loading will come later and have a much bigger
impact, but lazy-loading of themes does not have negligible impact.

CHANGELOG.md

@@ -17,6 +17,7 @@
 
 - Load cached assets as fast as integrated assets, see #1753 (@Enselic)
 - Greatly reduce startup time in loop-through mode, e.g. when redirecting output. Instead of *50 ms* - *100 ms*, startup takes *5 ms* - *10 ms*. See #1747 (@Enselic)
+- Load themes lazily to make bat start 25% faster when disregarding syntax load time. See #1969 (@Enselic)
 
 ## Other
 

src/assets.rs

@@ -4,7 +4,7 @@ use std::path::Path;
 
 use once_cell::unsync::OnceCell;
 
-use syntect::highlighting::{Theme, ThemeSet};
+use syntect::highlighting::Theme;
 use syntect::parsing::{SyntaxReference, SyntaxSet};
 
 use path_abs::PathAbs;
@@ -15,6 +15,8 @@ use crate::syntax_mapping::ignored_suffixes::IgnoredSuffixes;
 use crate::syntax_mapping::MappingTarget;
 use crate::{bat_warning, SyntaxMapping};
 
+use lazy_theme_set::LazyThemeSet;
+
 use serialized_syntax_set::*;
 
 #[cfg(feature = "build-assets")]
@@ -23,6 +25,7 @@ pub use crate::assets::build_assets::*;
 pub(crate) mod assets_metadata;
 #[cfg(feature = "build-assets")]
 mod build_assets;
+mod lazy_theme_set;
 mod serialized_syntax_set;
 
 #[derive(Debug)]
@@ -30,7 +33,7 @@ pub struct HighlightingAssets {
     syntax_set_cell: OnceCell<SyntaxSet>,
     serialized_syntax_set: SerializedSyntaxSet,
 
-    theme_set: ThemeSet,
+    theme_set: LazyThemeSet,
     fallback_theme: Option<&'static str>,
 }
 
@@ -43,11 +46,17 @@ pub struct SyntaxReferenceInSet<'a> {
 /// Compress for size of ~700 kB instead of ~4600 kB at the cost of ~30% longer deserialization time
 pub(crate) const COMPRESS_SYNTAXES: bool = true;
 
-/// Compress for size of ~20 kB instead of ~200 kB at the cost of ~30% longer deserialization time
-pub(crate) const COMPRESS_THEMES: bool = true;
+/// We don't want to compress our [LazyThemeSet] since the lazy-loaded themes
+/// within it are already compressed, and compressing another time just makes
+/// performance suffer
+pub(crate) const COMPRESS_THEMES: bool = false;
+
+/// Compress for size of ~40 kB instead of ~200 kB without much difference in
+/// performance due to lazy-loading
+pub(crate) const COMPRESS_LAZY_THEMES: bool = true;
 
 impl HighlightingAssets {
-    fn new(serialized_syntax_set: SerializedSyntaxSet, theme_set: ThemeSet) -> Self {
+    fn new(serialized_syntax_set: SerializedSyntaxSet, theme_set: LazyThemeSet) -> Self {
         HighlightingAssets {
             syntax_set_cell: OnceCell::new(),
             serialized_syntax_set,
@@ -95,12 +104,12 @@ impl HighlightingAssets {
         Ok(self.get_syntax_set()?.syntaxes())
     }
 
-    fn get_theme_set(&self) -> &ThemeSet {
+    fn get_theme_set(&self) -> &LazyThemeSet {
         &self.theme_set
     }
 
     pub fn themes(&self) -> impl Iterator<Item = &str> {
-        self.get_theme_set().themes.keys().map(|s| s.as_ref())
+        self.get_theme_set().themes()
     }
 
     /// Use [Self::get_syntax_for_path] instead
@@ -175,7 +184,7 @@ impl HighlightingAssets {
     }
 
     pub(crate) fn get_theme(&self, theme: &str) -> &Theme {
-        match self.get_theme_set().themes.get(theme) {
+        match self.get_theme_set().get(theme) {
             Some(theme) => theme,
             None => {
                 if theme == "ansi-light" || theme == "ansi-dark" {
@@ -185,8 +194,9 @@ impl HighlightingAssets {
                 if !theme.is_empty() {
                     bat_warning!("Unknown theme '{}', using default.", theme)
                 }
-                &self.get_theme_set().themes
-                    [self.fallback_theme.unwrap_or_else(|| Self::default_theme())]
+                self.get_theme_set()
+                    .get(self.fallback_theme.unwrap_or_else(|| Self::default_theme()))
+                    .expect("something is very wrong if the default theme is missing")
             }
         }
     }
@@ -291,7 +301,7 @@ pub(crate) fn get_serialized_integrated_syntaxset() -> &'static [u8] {
     include_bytes!("../assets/syntaxes.bin")
 }
 
-pub(crate) fn get_integrated_themeset() -> ThemeSet {
+pub(crate) fn get_integrated_themeset() -> LazyThemeSet {
     from_binary(include_bytes!("../assets/themes.bin"), COMPRESS_THEMES)
 }
 

src/assets/build_assets.rs

@@ -1,4 +1,6 @@
+use std::convert::TryInto;
 use std::path::Path;
+
 use syntect::highlighting::ThemeSet;
 use syntect::parsing::{SyntaxSet, SyntaxSetBuilder};
 
@@ -10,7 +12,7 @@ pub fn build(
     target_dir: &Path,
     current_version: &str,
 ) -> Result<()> {
-    let theme_set = build_theme_set(source_dir, include_integrated_assets);
+    let theme_set = build_theme_set(source_dir, include_integrated_assets)?;
 
     let syntax_set_builder = build_syntax_set_builder(source_dir, include_integrated_assets)?;
 
@@ -21,9 +23,9 @@ pub fn build(
     write_assets(&theme_set, &syntax_set, target_dir, current_version)
 }
 
-fn build_theme_set(source_dir: &Path, include_integrated_assets: bool) -> ThemeSet {
+fn build_theme_set(source_dir: &Path, include_integrated_assets: bool) -> Result<LazyThemeSet> {
     let mut theme_set = if include_integrated_assets {
-        crate::assets::get_integrated_themeset()
+        crate::assets::get_integrated_themeset().try_into()?
     } else {
         ThemeSet::new()
     };
@@ -45,7 +47,7 @@ fn build_theme_set(source_dir: &Path, include_integrated_assets: bool) -> ThemeS
         );
     }
 
-    theme_set
+    theme_set.try_into()
 }
 
 fn build_syntax_set_builder(
@@ -85,7 +87,7 @@ fn print_unlinked_contexts(syntax_set: &SyntaxSet) {
 }
 
 fn write_assets(
-    theme_set: &ThemeSet,
+    theme_set: &LazyThemeSet,
     syntax_set: &SyntaxSet,
     target_dir: &Path,
     current_version: &str,
@@ -114,7 +116,7 @@ fn write_assets(
     Ok(())
 }
 
-fn asset_to_contents<T: serde::Serialize>(
+pub(crate) fn asset_to_contents<T: serde::Serialize>(
     asset: &T,
     description: &str,
     compressed: bool,

src/assets/lazy_theme_set.rs

@@ -0,0 +1,104 @@
+use super::*;
+
+use std::collections::BTreeMap;
+use std::convert::TryFrom;
+
+use serde::Deserialize;
+use serde::Serialize;
+
+use once_cell::unsync::OnceCell;
+
+use syntect::highlighting::{Theme, ThemeSet};
+
+/// Same structure as a [`syntect::highlighting::ThemeSet`] but with themes
+/// stored in raw serialized form, and deserialized on demand.
+#[derive(Debug, Default, Serialize, Deserialize)]
+pub struct LazyThemeSet {
+    /// This is a [`BTreeMap`] because that's what [`syntect::highlighting::ThemeSet`] uses
+    themes: BTreeMap<String, LazyTheme>,
+}
+
+/// Stores raw serialized data for a theme with methods to lazily deserialize
+/// (load) the theme.
+#[derive(Debug, Serialize, Deserialize)]
+struct LazyTheme {
+    serialized: Vec<u8>,
+
+    #[serde(skip, default = "OnceCell::new")]
+    deserialized: OnceCell<syntect::highlighting::Theme>,
+}
+
+impl LazyThemeSet {
+    /// Lazily load the given theme
+    pub fn get(&self, name: &str) -> Option<&Theme> {
+        self.themes.get(name).and_then(|lazy_theme| {
+            lazy_theme
+                .deserialized
+                .get_or_try_init(|| lazy_theme.deserialize())
+                .ok()
+        })
+    }
+
+    /// Returns the name of all themes.
+    pub fn themes(&self) -> impl Iterator<Item = &str> {
+        self.themes.keys().map(|name| name.as_ref())
+    }
+}
+
+impl LazyTheme {
+    fn deserialize(&self) -> Result<Theme> {
+        asset_from_contents(
+            &self.serialized[..],
+            "lazy-loaded theme",
+            COMPRESS_LAZY_THEMES,
+        )
+    }
+}
+
+impl TryFrom<LazyThemeSet> for ThemeSet {
+    type Error = Error;
+
+    /// Since the user might want to add custom themes to bat, we need a way to
+    /// convert from a `LazyThemeSet` to a regular [`ThemeSet`] so that more
+    /// themes can be added. This function does that pretty straight-forward
+    /// conversion.
+    fn try_from(lazy_theme_set: LazyThemeSet) -> Result<Self> {
+        let mut theme_set = ThemeSet::default();
+
+        for (name, lazy_theme) in lazy_theme_set.themes {
+            theme_set.themes.insert(name, lazy_theme.deserialize()?);
+        }
+
+        Ok(theme_set)
+    }
+}
+
+#[cfg(feature = "build-assets")]
+impl TryFrom<ThemeSet> for LazyThemeSet {
+    type Error = Error;
+
+    /// To collect themes, a [`ThemeSet`] is needed. Once all desired themes
+    /// have been added, we need a way to convert that into [`LazyThemeSet`] so
+    /// that themes can be lazy-loaded later. This function does that
+    /// conversion.
+    fn try_from(theme_set: ThemeSet) -> Result<Self> {
+        let mut lazy_theme_set = LazyThemeSet::default();
+
+        for (name, theme) in theme_set.themes {
+            // All we have to do is to serialize the theme
+            let lazy_theme = LazyTheme {
+                serialized: crate::assets::build_assets::asset_to_contents(
+                    &theme,
+                    &format!("theme {}", name),
+                    COMPRESS_LAZY_THEMES,
+                )?,
+                deserialized: OnceCell::new(),
+            };
+
+            // Ok done, now we can add it
+            lazy_theme_set.themes.insert(name, lazy_theme);
+        }
+
+        Ok(lazy_theme_set)
+    }
+}