use std::collections::{BTreeMap, BTreeSet};
use rayon::iter::{IntoParallelRefIterator, ParallelIterator};
use crate::{
db::models::{Dependency, DownloadedBottle},
formulae::Formula,
};
#[derive(Debug)]
pub struct BottleForestSnapshot {
bottles: BTreeMap<i32, DownloadedBottle>,
dependencies: Vec<Dependency>,
}
#[derive(Debug)]
pub struct ActionBuilder<'a> {
snapshot: &'a BottleForestSnapshot,
bottles: BTreeSet<BottleRef<'a>>,
dependencies: BTreeSet<(Option<BottleRef<'a>>, BottleRef<'a>)>,
}
#[derive(Clone, Copy, Debug, Default, PartialEq, Eq, PartialOrd, Ord)]
struct BottleRef<'a> {
name: &'a str,
version: &'a str,
}
impl BottleForestSnapshot {
pub fn new() -> anyhow::Result<Self> {
let bottles = DownloadedBottle::get_all()?
.into_iter()
.map(|b| (b.id(), b))
.collect();
let dependencies = Dependency::get_all()?;
Ok(Self {
bottles,
dependencies,
})
}
}
impl<'a> ActionBuilder<'a> {
pub fn new(snapshot: &'a BottleForestSnapshot) -> Self {
let bottles = snapshot.bottles.values().map(BottleRef::from).collect();
let dependencies = snapshot
.dependencies
.iter()
.map(|dep| {
(
dep.dependent_id()
.map(|id| snapshot.bottles.get(&id).unwrap().into()),
snapshot.bottles.get(&dep.dependency_id()).unwrap().into(),
)
})
.collect();
Self {
snapshot,
bottles,
dependencies,
}
}
pub fn add_bottles(mut self, bottles: &[String]) -> anyhow::Result<Self> {
for name in bottles {
if self
.bottles
.range(BottleRef { name, version: "" }..)
.take_while(|b| b.name == name)
.count()
> 0
{
continue;
}
let formula = Formula::get(name)?;
self.bottles.insert(formula.into());
self.dependencies.insert((None, formula.into()));
}
Ok(self)
}
pub fn remove_all(mut self) -> Self {
self.bottles.clear();
self
}
pub fn remove_bottles(mut self, bottles: &'a [String]) -> anyhow::Result<Self> {
for alias in bottles {
let formula = Formula::get(alias);
let name = formula.as_ref().map_or(alias.as_str(), |f| &f.name);
let bottles_with_name = self
.bottles
.range(BottleRef { name, version: "" }..)
.take_while(|b| b.name == name)
.copied()
.collect::<Vec<BottleRef<'a>>>();
if bottles_with_name.is_empty() {
if formula.is_ok() {
anyhow::bail!("Could not remove {name} as it is not installed");
} else {
anyhow::bail!("No such formula {name}");
}
}
for bottle in bottles_with_name {
self.bottles.remove(&bottle);
}
}
Ok(self)
}
pub fn update(mut self) -> anyhow::Result<Self> {
let roots = self
.dependencies
.iter()
.filter(|(a, _)| a.is_none())
.map(|(_, b)| Formula::get_exact(b.name))
.collect::<Result<Vec<_>, _>>()?;
self.bottles = roots.iter().copied().map(BottleRef::from).collect();
self.dependencies = roots
.iter()
.copied()
.map(|f| (None, BottleRef::from(f)))
.collect();
Ok(self)
}
pub fn run(mut self) -> anyhow::Result<()> {
self.fix_dependencies()?;
let (to_add, to_remove) = diff_bottles(
&self
.snapshot
.bottles
.values()
.map(BottleRef::from)
.collect(),
&self.bottles,
);
anyhow::ensure!(
!to_add.is_empty() || !to_remove.is_empty(),
"No bottles to add or remove",
);
let downloaded_bottles = to_add
.par_iter()
.map(|bottle_ref| {
let formula = Formula::get_exact(bottle_ref.name)?;
anyhow::ensure!(
formula.versions.stable == bottle_ref.version,
"Attempted to install an unavailable version of {}",
bottle_ref.name,
);
anyhow::ensure!(
formula.versions.bottle,
"Formula {:?} does not have a corresponding bottle",
formula.name,
);
let bottle = formula.download_bottle()?;
Ok(bottle)
})
.collect::<anyhow::Result<Vec<_>>>()?;
downloaded_bottles
.par_iter()
.map(|bottle| {
bottle.link()?;
Ok(())
})
.collect::<anyhow::Result<Vec<()>>>()?;
let bottles_by_ref = self
.snapshot
.bottles
.values()
.chain(&downloaded_bottles)
.map(|b| (BottleRef::from(b), b))
.collect::<BTreeMap<_, _>>();
Dependency::replace_all(
self.dependencies
.iter()
.map(|(a, b)| (a.map(|a| bottles_by_ref[&a]), bottles_by_ref[b])),
)?;
to_remove
.par_iter()
.map(|bottle_ref| {
let bottle = bottles_by_ref[bottle_ref];
bottle.unlink()?;
bottle.remove()?;
Ok(())
})
.collect::<anyhow::Result<Vec<()>>>()?;
Ok(())
}
fn get_bottle(&self, name: &'a str) -> Option<BottleRef<'a>> {
self.bottles
.range(BottleRef { name, version: "" }..)
.take_while(|b| b.name == name)
.cloned()
.next()
}
fn get_dependencies(&self, bottle_ref: BottleRef<'a>) -> impl Iterator<Item = BottleRef<'a>> {
self.dependencies
.range(
(
Some(bottle_ref),
BottleRef {
name: "",
version: "",
},
)..,
)
.take_while(move |(a, _)| a == &Some(bottle_ref))
.map(|&(_, b)| b)
}
fn fix_dependencies(&mut self) -> anyhow::Result<()> {
self.add_dependencies()?;
self.remove_orphans();
Ok(())
}
fn add_dependencies(&mut self) -> Result<(), anyhow::Error> {
let mut stack = Vec::new();
for bottle in self.bottles.iter() {
let Ok(formula) = Formula::get_exact(bottle.name) else {
continue;
};
if formula.versions.stable != bottle.version {
continue;
}
stack.push(formula);
}
while let Some(formula) = stack.pop() {
let bottle_ref = BottleRef::from(formula);
for dependency_name in &formula.dependencies {
if let Some(dependency_ref) = self.get_bottle(dependency_name) {
self.dependencies.insert((Some(bottle_ref), dependency_ref));
continue;
}
let dependency = Formula::get_exact(dependency_name)?;
let dependency_ref = BottleRef::from(dependency);
self.bottles.insert(dependency_ref);
self.dependencies.insert((Some(bottle_ref), dependency_ref));
stack.push(dependency);
}
}
Ok(())
}
fn remove_orphans(&mut self) {
let mut ref_counts = self
.bottles
.iter()
.map(|&b| (b, 0))
.collect::<BTreeMap<_, _>>();
self.dependencies.retain(|&(a, b)| {
if let Some(bottle) = a {
if !self.bottles.contains(&bottle) {
return false;
}
}
if !self.bottles.contains(&b) {
return false;
}
let ref_count = ref_counts.get_mut(&b).unwrap();
*ref_count += 1;
true
});
let mut stack = Vec::new();
for bottle in self.bottles.iter() {
if ref_counts[bottle] == 0 {
stack.push(*bottle);
}
}
while let Some(bottle) = stack.pop() {
self.bottles.remove(&bottle);
for dependency in self.get_dependencies(bottle).collect::<Vec<_>>() {
self.dependencies.remove(&(Some(bottle), dependency));
let ref_count = ref_counts.get_mut(&dependency).unwrap();
*ref_count -= 1;
if *ref_count == 0 {
stack.push(dependency);
}
}
}
}
}
fn diff_bottles<'a>(
before: &BTreeSet<BottleRef<'a>>,
after: &BTreeSet<BottleRef<'a>>,
) -> (BTreeSet<BottleRef<'a>>, BTreeSet<BottleRef<'a>>) {
let added = after.difference(before).cloned().collect();
let removed = before.difference(after).cloned().collect();
(added, removed)
}
impl<'a> From<&'a DownloadedBottle> for BottleRef<'a> {
fn from(bottle: &'a DownloadedBottle) -> Self {
Self {
name: bottle.name(),
version: bottle.version(),
}
}
}
impl<'a> From<&'a Formula> for BottleRef<'a> {
fn from(formula: &'a Formula) -> Self {
Self {
name: &formula.name,
version: &formula.versions.stable,
}
}
}