use super::{Graph, Node, NodeLayer};
const TOKENS_PER_LINE: usize = 10;
#[derive(Debug, Clone, PartialEq, Eq, serde::Serialize, serde::Deserialize)]
#[serde(rename_all = "snake_case")]
pub enum ContextMode {
Map,
Lite,
Compact,
Full,
FullExtended,
Custom,
Preset(String),
}
impl ContextMode {
pub fn name(&self) -> &'static str {
match self {
Self::Map => "map",
Self::Lite => "lite",
Self::Compact => "compact",
Self::Full => "full",
Self::FullExtended => "full_extended",
Self::Custom => "custom",
Self::Preset(_) => "preset",
}
}
}
#[derive(Debug, Clone, PartialEq, Eq, serde::Serialize, serde::Deserialize)]
pub struct ContextModeSize {
pub mode: String,
pub nodes: usize,
pub tokens: usize,
}
#[derive(Debug, Clone, PartialEq, Eq, serde::Serialize, serde::Deserialize)]
pub struct ContextLayerSummary {
pub layer: NodeLayer,
pub nodes: usize,
pub tokens: usize,
pub files: usize,
}
#[derive(Debug, Clone, Default, PartialEq, Eq, serde::Serialize, serde::Deserialize)]
pub struct ContextSourceSummary {
pub nodes: usize,
pub tokens: usize,
pub files: usize,
}
#[derive(Debug, Clone, PartialEq, Eq, serde::Serialize, serde::Deserialize)]
pub struct ContextSummary {
pub mode: ContextMode,
pub included: Vec<String>,
pub estimated_tokens: usize,
pub production: ContextSourceSummary,
pub tests: ContextSourceSummary,
pub examples: ContextSourceSummary,
pub file_partition_complete: bool,
pub production_files_with_inline_tests: usize,
pub inline_test_ranges: usize,
pub inline_test_lines: usize,
pub layers: Vec<ContextLayerSummary>,
}
pub struct ContextComposer {
graph: Graph,
mode: ContextMode,
included: Vec<String>,
}
impl ContextComposer {
pub fn new(graph: Graph) -> Self {
Self {
graph,
mode: ContextMode::Lite,
included: Vec::new(),
}
}
pub fn set_mode(&mut self, mode: ContextMode) {
self.included = self.included_for(&mode);
self.mode = mode;
}
pub fn set_custom(&mut self, ids: Vec<String>) {
let known: std::collections::HashSet<&str> =
self.graph.nodes.iter().map(|n| n.id.as_str()).collect();
let mut seen = std::collections::HashSet::new();
self.included = ids
.into_iter()
.filter(|id| known.contains(id.as_str()) && seen.insert(id.clone()))
.collect();
self.mode = ContextMode::Custom;
}
fn included_for(&self, mode: &ContextMode) -> Vec<String> {
match mode {
ContextMode::Map | ContextMode::Lite | ContextMode::Compact | ContextMode::Full => self
.graph
.nodes
.iter()
.filter(|node| node.layer == NodeLayer::Code)
.map(|n| n.id.clone())
.collect(),
ContextMode::FullExtended => self
.graph
.nodes
.iter()
.filter(|node| matches!(node.layer, NodeLayer::Code | NodeLayer::Test))
.map(|node| node.id.clone())
.collect(),
ContextMode::Preset(name) => vec![name.clone()],
ContextMode::Custom => self.included.clone(),
}
}
pub fn mode_sizes(&self) -> Vec<ContextModeSize> {
[
ContextMode::Lite,
ContextMode::Compact,
ContextMode::Full,
ContextMode::FullExtended,
]
.into_iter()
.map(|mode| {
let included = self.included_for(&mode);
let tokens = self
.graph
.nodes
.iter()
.filter(|node| included.contains(&node.id))
.map(|node| estimate_context_node_tokens(node, &mode))
.sum();
ContextModeSize {
mode: mode.name().to_string(),
nodes: included.len(),
tokens,
}
})
.collect()
}
pub fn mode_name(&self) -> &'static str {
self.mode.name()
}
pub fn mode(&self) -> &ContextMode {
&self.mode
}
pub fn estimate_tokens(&self) -> usize {
self.graph
.nodes
.iter()
.filter(|n| self.included.contains(&n.id))
.map(|node| estimate_context_node_tokens(node, &self.mode))
.sum()
}
pub fn included_nodes(&self) -> Vec<String> {
self.included.clone()
}
pub fn layer_summaries(&self) -> Vec<ContextLayerSummary> {
[
NodeLayer::Code,
NodeLayer::Test,
NodeLayer::Concept,
NodeLayer::Preset,
]
.into_iter()
.filter_map(|layer| {
let nodes: Vec<_> = self
.graph
.nodes
.iter()
.filter(|node| node.layer == layer && self.included.contains(&node.id))
.collect();
(!nodes.is_empty()).then(|| ContextLayerSummary {
layer,
nodes: nodes.len(),
tokens: nodes
.iter()
.map(|node| estimate_context_node_tokens(node, &self.mode))
.sum(),
files: nodes.iter().map(|node| node.files).sum(),
})
})
.collect()
}
pub fn summary(&self) -> ContextSummary {
let mut production = ContextSourceSummary::default();
let mut tests = ContextSourceSummary::default();
let mut examples = ContextSourceSummary::default();
let mut production_files_with_inline_tests = 0usize;
let mut inline_test_ranges = 0usize;
let mut inline_test_lines = 0usize;
for node in self
.graph
.nodes
.iter()
.filter(|node| self.included.contains(&node.id))
{
let bucket = if node.layer == NodeLayer::Code {
&mut production
} else if node.layer == NodeLayer::Test && node.id.starts_with("example::") {
&mut examples
} else if node.layer == NodeLayer::Test {
&mut tests
} else {
continue;
};
bucket.nodes += 1;
bucket.tokens += estimate_context_node_tokens(node, &self.mode);
bucket.files += node.files;
if node.layer == NodeLayer::Code && node.inline_test_ranges > 0 {
production_files_with_inline_tests += 1;
inline_test_ranges += node.inline_test_ranges;
inline_test_lines += node.inline_test_lines;
}
}
ContextSummary {
mode: self.mode.clone(),
included: self.included.clone(),
estimated_tokens: self.estimate_tokens(),
production,
tests,
examples,
file_partition_complete: production_files_with_inline_tests == 0,
production_files_with_inline_tests,
inline_test_ranges,
inline_test_lines,
layers: self.layer_summaries(),
}
}
}
fn estimate_node_tokens(tokens: usize, lines: usize) -> usize {
if tokens > 0 {
tokens
} else {
(lines * TOKENS_PER_LINE).max(1)
}
}
const SIGNATURE_FRACTION: f64 = 0.18;
const COMMENT_STRIPPED_FRACTION: f64 = 0.82;
fn estimate_context_node_tokens(node: &Node, mode: &ContextMode) -> usize {
let total = estimate_node_tokens(node.tokens, node.lines);
let code_tokens = total.saturating_sub(node.inline_test_tokens);
match mode {
ContextMode::Map => 0,
ContextMode::Lite | ContextMode::Custom => {
((code_tokens as f64) * SIGNATURE_FRACTION).round() as usize
}
ContextMode::Compact => ((code_tokens as f64) * COMMENT_STRIPPED_FRACTION).round() as usize,
ContextMode::Full => code_tokens,
_ => total,
}
}
#[cfg(test)]
#[path = "../../tests/unit/evolve/context_custom_test.rs"]
mod context_custom_test;