use std::collections::BTreeSet;
use std::path::{Component, Path, PathBuf};
use crate::pipeline::Flags;
pub const MAX_INCLUDE_DEPTH: usize = 8;
pub const MAX_INCLUDE_EXPANSIONS: usize = 1_000;
pub trait PromptSource {
fn read(&self, path: &Path) -> Result<String, PromptError>;
}
#[derive(Debug, Clone)]
pub struct PromptDir {
root: PathBuf,
}
impl PromptDir {
#[must_use]
pub fn new(root: impl Into<PathBuf>) -> Self {
Self { root: root.into() }
}
#[must_use]
pub fn root(&self) -> &Path {
&self.root
}
}
impl PromptSource for PromptDir {
fn read(&self, path: &Path) -> Result<String, PromptError> {
let relative = confine(path)?;
let full = self.root.join(&relative);
if let Some(escape) = self.escapes(&full) {
return Err(PromptError::Escapes { path: escape });
}
match std::fs::read_to_string(&full) {
Ok(text) => Ok(text),
Err(error) if error.kind() == std::io::ErrorKind::NotFound => {
Err(PromptError::Missing { path: relative })
}
Err(error) => Err(PromptError::Unreadable {
path: relative,
reason: error.to_string(),
}),
}
}
}
impl PromptDir {
fn escapes(&self, full: &Path) -> Option<PathBuf> {
let resolved = std::fs::canonicalize(full).ok()?;
let root = std::fs::canonicalize(&self.root).ok()?;
(!resolved.starts_with(&root)).then(|| {
full.strip_prefix(&self.root).unwrap_or(full).to_path_buf()
})
}
}
#[derive(Debug, Clone, Default)]
pub struct PromptMap {
files: std::collections::BTreeMap<PathBuf, String>,
}
impl PromptMap {
#[must_use]
pub fn new() -> Self {
Self::default()
}
#[must_use]
pub fn with(mut self, path: impl Into<PathBuf>, text: impl Into<String>) -> Self {
self.files.insert(path.into(), text.into());
self
}
}
impl PromptSource for PromptMap {
fn read(&self, path: &Path) -> Result<String, PromptError> {
let relative = confine(path)?;
self.files
.get(&relative)
.cloned()
.ok_or(PromptError::Missing { path: relative })
}
}
pub fn resolve(
source: &dyn PromptSource,
entry: impl AsRef<Path>,
flags: &Flags,
) -> Result<String, PromptError> {
let mut stack = Vec::new();
let mut output = String::new();
expand(
source,
&confine(entry.as_ref())?,
flags,
&mut stack,
&mut output,
)?;
Ok(output)
}
pub fn referenced_flags(
source: &dyn PromptSource,
entry: impl AsRef<Path>,
) -> Result<BTreeSet<String>, PromptError> {
let mut found = BTreeSet::new();
let mut stack = Vec::new();
let mut budget = MAX_INCLUDE_EXPANSIONS;
walk(
source,
&confine(entry.as_ref())?,
&mut stack,
&mut budget,
&mut found,
)?;
Ok(found)
}
fn walk(
source: &dyn PromptSource,
path: &Path,
stack: &mut Vec<PathBuf>,
budget: &mut usize,
found: &mut BTreeSet<String>,
) -> Result<(), PromptError> {
if stack.len() >= MAX_INCLUDE_DEPTH {
return Err(PromptError::TooDeep {
path: path.to_path_buf(),
limit: MAX_INCLUDE_DEPTH,
});
}
if stack.iter().any(|seen| seen == path) {
return Err(PromptError::Cycle {
path: path.to_path_buf(),
});
}
*budget = budget.checked_sub(1).ok_or_else(|| PromptError::TooWide {
path: path.to_path_buf(),
limit: MAX_INCLUDE_EXPANSIONS,
})?;
let text = source.read(path)?;
stack.push(path.to_path_buf());
for (number, line) in text.lines().enumerate() {
let Some(directive) = Directive::parse(line, path, number + 1)? else {
continue;
};
if let Some(condition) = &directive.condition {
found.insert(condition.flag.clone());
}
let target = confine(&resolve_relative(path, &directive.path))?;
walk(source, &target, stack, budget, found)?;
}
stack.pop();
Ok(())
}
fn expand(
source: &dyn PromptSource,
path: &Path,
flags: &Flags,
stack: &mut Vec<PathBuf>,
output: &mut String,
) -> Result<(), PromptError> {
if stack.len() >= MAX_INCLUDE_DEPTH {
return Err(PromptError::TooDeep {
path: path.to_path_buf(),
limit: MAX_INCLUDE_DEPTH,
});
}
if stack.iter().any(|seen| seen == path) {
return Err(PromptError::Cycle {
path: path.to_path_buf(),
});
}
let text = source.read(path)?;
stack.push(path.to_path_buf());
for (number, line) in text.lines().enumerate() {
match Directive::parse(line, path, number + 1)? {
None => {
output.push_str(line);
output.push('\n');
}
Some(directive) => {
let include = match &directive.condition {
None => true,
Some(condition) => {
let value =
flags
.get(&condition.flag)
.ok_or_else(|| PromptError::UnknownFlag {
flag: condition.flag.clone(),
path: path.to_path_buf(),
line: number + 1,
})?;
value != condition.negated
}
};
if include {
let target = confine(&resolve_relative(path, &directive.path))?;
expand(source, &target, flags, stack, output)?;
}
}
}
}
stack.pop();
Ok(())
}
#[derive(Debug, Clone, PartialEq, Eq)]
struct Directive {
condition: Option<Condition>,
path: PathBuf,
}
#[derive(Debug, Clone, PartialEq, Eq)]
struct Condition {
flag: String,
negated: bool,
}
impl Directive {
fn parse(line: &str, path: &Path, number: usize) -> Result<Option<Self>, PromptError> {
let trimmed = line.trim();
let Some(rest) = trimmed.strip_prefix("@include") else {
return Ok(None);
};
let malformed = || PromptError::MalformedDirective {
path: path.to_path_buf(),
line: number,
text: trimmed.to_owned(),
};
let (condition, remainder) = if let Some(after) = rest.strip_prefix('(') {
let (inside, after) = after.split_once(')').ok_or_else(malformed)?;
(Some(parse_condition(inside).ok_or_else(malformed)?), after)
} else if rest.starts_with(char::is_whitespace) {
(None, rest)
} else {
return Err(malformed());
};
let target = remainder.trim().trim_matches('"').trim();
if target.is_empty() {
return Err(malformed());
}
Ok(Some(Self {
condition,
path: PathBuf::from(target),
}))
}
}
fn parse_condition(inside: &str) -> Option<Condition> {
let trimmed = inside.trim();
let (negated, name) = match trimmed.strip_prefix('!') {
Some(rest) => (true, rest.trim()),
None => (false, trimmed),
};
let mut chars = name.chars();
let first = chars.next()?;
if !(first.is_ascii_alphabetic() || first == '_') {
return None;
}
if !chars.all(|c| c.is_ascii_alphanumeric() || c == '_') {
return None;
}
Some(Condition {
flag: name.to_owned(),
negated,
})
}
fn resolve_relative(including: &Path, target: &Path) -> PathBuf {
match including.parent() {
Some(parent) if !parent.as_os_str().is_empty() => parent.join(target),
_ => target.to_path_buf(),
}
}
fn confine(path: &Path) -> Result<PathBuf, PromptError> {
let escapes = || PromptError::Escapes {
path: path.to_path_buf(),
};
let mut clean = PathBuf::new();
for component in path.components() {
match component {
Component::Normal(part) => clean.push(part),
Component::CurDir => {}
Component::ParentDir => {
if !clean.pop() {
return Err(escapes());
}
}
Component::RootDir | Component::Prefix(_) => return Err(escapes()),
}
}
if clean.as_os_str().is_empty() {
return Err(escapes());
}
Ok(clean)
}
#[derive(Debug, Clone, PartialEq, Eq, thiserror::Error)]
pub enum PromptError {
#[error("prompt file `{path}` does not exist")]
Missing {
path: PathBuf,
},
#[error("prompt file `{path}` could not be read: {reason}")]
Unreadable {
path: PathBuf,
reason: String,
},
#[error("prompt path `{path}` leaves the prompt directory")]
Escapes {
path: PathBuf,
},
#[error("`{path}` line {line}: could not read `{text}` as an @include directive")]
MalformedDirective {
path: PathBuf,
line: usize,
text: String,
},
#[error("`{path}` line {line}: flag `{flag}` is not declared by any pipeline")]
UnknownFlag {
flag: String,
path: PathBuf,
line: usize,
},
#[error("prompt file `{path}` includes itself")]
Cycle {
path: PathBuf,
},
#[error("prompt file `{path}` nests includes more than {limit} deep")]
TooDeep {
path: PathBuf,
limit: usize,
},
#[error("prompt composition reached `{path}` after more than {limit} expansions")]
TooWide {
path: PathBuf,
limit: usize,
},
}