use std::collections::HashMap;
use std::fmt;
#[derive(Clone, Debug)]
pub struct Template {
parts: Vec<TemplatePart>,
estimated_size: usize,
}
#[derive(Clone, Debug)]
enum TemplatePart {
Literal(String),
Field(String),
}
#[derive(Debug)]
pub struct TemplateError {
pub reason: String,
}
impl fmt::Display for TemplateError {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
write!(f, "invalid template: {}", self.reason)
}
}
impl std::error::Error for TemplateError {}
impl Template {
pub fn compile(template: &str) -> Result<Template, TemplateError> {
let mut parts = Vec::new();
let mut literal = String::new();
let mut estimated_size = 0;
let bytes = template.as_bytes();
let len = bytes.len();
let mut i = 0;
while i < len {
if bytes[i] == b'{' {
if i + 1 < len && bytes[i + 1] == b'{' {
literal.push('{');
i += 2;
continue;
}
if let Some(close) = template[i + 1..].find('}') {
let field_name = &template[i + 1..i + 1 + close];
if field_name.is_empty() {
return Err(TemplateError {
reason: "empty field name at position".to_string(),
});
}
if !literal.is_empty() {
estimated_size += literal.len();
parts.push(TemplatePart::Literal(std::mem::take(&mut literal)));
}
estimated_size += 16;
parts.push(TemplatePart::Field(field_name.to_string()));
i += 1 + close + 1; } else {
literal.push('{');
i += 1;
}
} else if bytes[i] == b'}' && i + 1 < len && bytes[i + 1] == b'}' {
literal.push('}');
i += 2;
} else {
literal.push(bytes[i] as char);
i += 1;
}
}
if !literal.is_empty() {
estimated_size += literal.len();
parts.push(TemplatePart::Literal(literal));
}
Ok(Template {
parts,
estimated_size,
})
}
#[inline]
pub fn render<'a>(&self, mut lookup: impl FnMut(&str) -> &'a str) -> String {
let mut output = String::with_capacity(self.estimated_size);
for part in &self.parts {
match part {
TemplatePart::Literal(s) => output.push_str(s),
TemplatePart::Field(name) => output.push_str(lookup(name)),
}
}
output
}
#[inline]
pub fn write<W, L>(&self, wtr: &mut W, mut lookup: L) -> std::io::Result<()>
where
W: std::io::Write + ?Sized,
L: FnMut(&mut W, &str) -> std::io::Result<()>,
{
for part in &self.parts {
match part {
TemplatePart::Literal(s) => wtr.write_all(s.as_bytes())?,
TemplatePart::Field(f) => {
lookup(wtr, f)?;
}
}
}
Ok(())
}
#[inline]
#[must_use]
pub fn render_with_map(&self, values: &HashMap<String, String>) -> String {
self.render(move |name| values.get(name).map_or("", |s| s.as_str()))
}
#[must_use]
pub fn fields(&self) -> Vec<&str> {
self.parts
.iter()
.filter_map(|part| match part {
TemplatePart::Field(name) => Some(name.as_str()),
TemplatePart::Literal(_) => None,
})
.collect()
}
}
impl fmt::Display for Template {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
for part in &self.parts {
match part {
TemplatePart::Literal(s) => write!(f, "{s}")?,
TemplatePart::Field(name) => write!(f, "{{{name}}}")?,
}
}
Ok(())
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn basic_template() {
let t = Template::compile("<{ip}|{country}>").unwrap();
let mut values = HashMap::new();
values.insert("ip".to_string(), "1.2.3.4".to_string());
values.insert("country".to_string(), "US".to_string());
assert_eq!(t.render_with_map(&values), "<1.2.3.4|US>");
}
#[test]
fn multiple_fields() {
let t = Template::compile("{a}-{b}-{c}").unwrap();
let result = t.render(|name| match name {
"a" => "1",
"b" => "2",
"c" => "3",
_ => "",
});
assert_eq!(result, "1-2-3");
}
#[test]
fn empty_template() {
let t = Template::compile("").unwrap();
assert_eq!(t.render(|_| "unused"), "");
}
#[test]
fn all_literal() {
let t = Template::compile("no fields here").unwrap();
assert_eq!(t.render(|_| "unused"), "no fields here");
}
#[test]
fn escaped_braces() {
let t = Template::compile("{{literal}} and {field}").unwrap();
assert_eq!(t.render(|_| "val"), "{literal} and val");
}
#[test]
fn escaped_closing_brace() {
let t = Template::compile("value is }}done").unwrap();
assert_eq!(t.render(|_| ""), "value is }done");
}
#[test]
fn no_double_substitution() {
let t = Template::compile("{a} and {b}").unwrap();
let result = t.render(|name| match name {
"a" => "{b}",
"b" => "real_b",
_ => "",
});
assert_eq!(result, "{b} and real_b");
}
#[test]
fn unknown_fields_empty() {
let t = Template::compile("{known} {unknown}").unwrap();
let mut values = HashMap::new();
values.insert("known".to_string(), "yes".to_string());
assert_eq!(t.render_with_map(&values), "yes ");
}
#[test]
fn fields_method() {
let t = Template::compile("{ip}|{asnnum}_{asnorg}|{country_iso}").unwrap();
assert_eq!(t.fields(), vec!["ip", "asnnum", "asnorg", "country_iso"]);
}
#[test]
fn unclosed_brace_is_literal() {
let t = Template::compile("value is {unclosed").unwrap();
assert_eq!(t.render(|_| ""), "value is {unclosed");
}
#[test]
fn display_roundtrip() {
let template_str = "<{ip}|AS{asnnum}_{asnorg}|{country_iso}|{city}>";
let t = Template::compile(template_str).unwrap();
assert_eq!(t.to_string(), template_str);
}
#[test]
fn empty_field_name_is_error() {
assert!(Template::compile("{}").is_err());
}
#[test]
fn geoipsed_default_template() {
let t = Template::compile("<{ip}|AS{asnnum}_{asnorg}|{country_iso}|{city}>").unwrap();
let result = t.render(|name| match name {
"ip" => "93.184.216.34",
"asnnum" => "15133",
"asnorg" => "EDGECAST",
"country_iso" => "US",
"city" => "Los_Angeles",
_ => "",
});
assert_eq!(result, "<93.184.216.34|AS15133_EDGECAST|US|Los_Angeles>");
}
}