use std::io::IsTerminal;
use bstr::ByteSlice;
use compact_genome::{
implementation::{alphabets::dna_alphabet_or_n::DnaAlphabetOrN, vec_sequence::VectorGenome},
interface::sequence::{GenomeSequence, OwnedGenomeSequence as _},
};
use lib_tsalign::{
a_star_aligner::{
alignment_result::AlignmentResult,
template_switch_distance::{
AlignmentType,
context::DynamicStrategies,
strategies::{
AlignmentStrategySelection,
allow_ts_14_out_of_range::{
AdditionalExplicitTSMStartsAndEnds, Ts14OutOfRangeStrategy,
},
chaining::NoChainingStrategy,
descendant::{
AnyTemplateSwitchDescendantStrategy, OnlyEqualTemplateSwitchDescendantStrategy,
TemplateSwitchDescendantStrategy,
},
node_ord::AntiDiagonalNodeOrdStrategy,
primary_match::AllowPrimaryMatchStrategy,
primary_range::NoPrunePrimaryRangeStrategy,
secondary_deletion::AllowSecondaryDeletionStrategy,
shortcut::NoShortcutStrategy,
template_switch_count::{
MaxTemplateSwitchCountStrategy, NoTemplateSwitchCountStrategy,
TemplateSwitchCountStrategy,
},
template_switch_min_length::{
LookaheadTemplateSwitchMinLengthStrategy, NoTemplateSwitchMinLengthStrategy,
PreprocessedLookaheadTemplateSwitchMinLengthStrategy,
PreprocessedTemplateSwitchMinLengthStrategy, TemplateSwitchMinLengthStrategy,
},
template_switch_total_length::MaxTemplateSwitchTotalLengthStrategy,
},
},
template_switch_distance_a_star_align,
},
config::TemplateSwitchConfig,
costs::U64Cost,
};
use rlimit::Resource;
use serde::{Deserialize, Serialize};
use tracing::{error, info_span, instrument};
use crate::{
VerbositySelector,
common::{
aligner::{
AlignerSelector, AlignmentQuery,
cli::MLSSelector,
result::{self, AlignmentFailure, TwitcherAlignmentResult},
},
coords::GenomeRegion,
},
};
pub enum RefOrOwned<'a, T: 'a> {
Ref(&'a T),
Owned(T),
}
impl<T> std::ops::Deref for RefOrOwned<'_, T> {
type Target = T;
fn deref(&self) -> &Self::Target {
match self {
RefOrOwned::Ref(r) => r,
RefOrOwned::Owned(o) => o,
}
}
}
impl<T: Serialize> Serialize for RefOrOwned<'_, T> {
fn serialize<S: serde::Serializer>(&self, serializer: S) -> Result<S::Ok, S::Error> {
match self {
RefOrOwned::Ref(r) => r.serialize(serializer),
RefOrOwned::Owned(o) => o.serialize(serializer),
}
}
}
impl<'de, T: Deserialize<'de>> Deserialize<'de> for RefOrOwned<'_, T> {
fn deserialize<D: serde::Deserializer<'de>>(deserializer: D) -> Result<Self, D::Error> {
T::deserialize(deserializer).map(RefOrOwned::Owned)
}
}
impl<T> From<T> for RefOrOwned<'_, T> {
fn from(value: T) -> Self {
Self::Owned(value)
}
}
impl<'a, T> From<&'a T> for RefOrOwned<'a, T> {
fn from(value: &'a T) -> Self {
Self::Ref(value)
}
}
#[derive(Serialize, Deserialize)]
pub struct WorkerQuery<'a> {
pub aligner: RefOrOwned<'a, AlignerSelector>,
pub log_level: VerbositySelector,
pub memory: usize,
pub query: AlignmentQuery,
pub metadata: WorkerQueryMetadata,
}
#[derive(Serialize, Deserialize)]
pub struct WorkerQueryMetadata {
pub cluster_region: GenomeRegion,
}
#[instrument]
pub fn run_fallible() -> anyhow::Result<()> {
let wq = read_from_stdin()?;
crate::setup_tracing(wq.log_level);
let span = info_span!("Aligning sequences", pos = %wq.metadata.cluster_region).entered();
let bytes = wq.memory as u64;
rlimit::setrlimit(Resource::AS, bytes, bytes)
.inspect_err(|e| error!("Can't set resource limit: {e}"))?;
let result = if let Err(e) = verify_query(&wq) {
TwitcherAlignmentResult::Err(AlignmentFailure::error(&e))
} else {
let WorkerQuery {
aligner,
memory,
query,
..
} = wq;
match &*aligner {
AlignerSelector::AStar {
costs,
min_length_strategy,
allow_mixed_descendants,
no_ts,
} => {
let with_ts = align_select_min_length_strategy(AStarQueryData {
memory,
query,
costs,
min_length_strategy,
allow_mixed_descendants: *allow_mixed_descendants,
no_ts: *no_ts,
});
result::from_tsalign(with_ts)
}
AlignerSelector::Fpa(four_point_aligner) => four_point_aligner.align(
query.sequences.reference,
query.sequences.query,
query.ranges,
),
}
};
span.exit();
let _span = info_span!("Reporting alignment", pos = %wq.metadata.cluster_region).entered();
let mut out = std::io::stdout().lock();
rmp_serde::encode::write(&mut out, &result)?;
std::mem::drop(out);
Ok(())
}
fn verify_query(wq: &WorkerQuery) -> anyhow::Result<()> {
let dnan = b"ACGTN";
for (s, bs) in [
("reference", &wq.query.sequences.reference),
("query", &wq.query.sequences.query),
] {
if let Some(c) = bs.find_not_byteset(dnan) {
anyhow::bail!(
"{s} sequence contains an unsupported ASCII character: {}",
bs[c..=c].as_bstr()
)
}
}
Ok(())
}
fn read_from_stdin() -> anyhow::Result<WorkerQuery<'static>> {
let mut stdin = std::io::stdin().lock();
if stdin.is_terminal() {
eprintln!(
"The `worker` subcommand is not designed to be called directly but only as a subroutine of twitcher. You will most likely not get the expected functionality."
);
}
let wq = rmp_serde::from_read(&mut stdin)?;
Ok(wq)
}
struct AStarQueryData<'a> {
memory: usize,
query: AlignmentQuery,
costs: &'a TemplateSwitchConfig<DnaAlphabetOrN, U64Cost>,
min_length_strategy: &'a MLSSelector,
allow_mixed_descendants: bool,
no_ts: bool,
}
fn align_select_min_length_strategy(
data: AStarQueryData,
) -> AlignmentResult<AlignmentType, U64Cost> {
match data.min_length_strategy {
MLSSelector::None => {
align_select_descendant_strategy::<NoTemplateSwitchMinLengthStrategy<U64Cost>>(data)
}
MLSSelector::Lookahead => align_select_descendant_strategy::<
LookaheadTemplateSwitchMinLengthStrategy<U64Cost>,
>(data),
MLSSelector::PreprocessFilter => align_select_descendant_strategy::<
PreprocessedTemplateSwitchMinLengthStrategy<true, U64Cost>,
>(data),
MLSSelector::PreprocessPrice => align_select_descendant_strategy::<
PreprocessedTemplateSwitchMinLengthStrategy<false, U64Cost>,
>(data),
MLSSelector::PreprocessLookahead => align_select_descendant_strategy::<
PreprocessedLookaheadTemplateSwitchMinLengthStrategy<U64Cost>,
>(data),
}
}
fn align_select_descendant_strategy<ML: TemplateSwitchMinLengthStrategy<U64Cost>>(
data: AStarQueryData,
) -> AlignmentResult<AlignmentType, U64Cost> {
if data.allow_mixed_descendants {
align_select_no_ts_strategy::<ML, AnyTemplateSwitchDescendantStrategy>(data)
} else {
align_select_no_ts_strategy::<ML, OnlyEqualTemplateSwitchDescendantStrategy>(data)
}
}
fn align_select_no_ts_strategy<
ML: TemplateSwitchMinLengthStrategy<U64Cost>,
DS: TemplateSwitchDescendantStrategy,
>(
data: AStarQueryData,
) -> AlignmentResult<AlignmentType, U64Cost> {
if data.no_ts {
align_call::<ML, DS, MaxTemplateSwitchCountStrategy>(data, 0)
} else {
align_call::<ML, DS, NoTemplateSwitchCountStrategy>(data, ())
}
}
fn align_call<
ML: TemplateSwitchMinLengthStrategy<U64Cost>,
DS: TemplateSwitchDescendantStrategy,
TS: TemplateSwitchCountStrategy,
>(
data: AStarQueryData,
count_mem: TS::Memory,
) -> AlignmentResult<AlignmentType, U64Cost> {
let reference = VectorGenome::<DnaAlphabetOrN>::from_iter_u8(
data.query.sequences.reference.iter().copied(),
)
.unwrap();
let query =
VectorGenome::<DnaAlphabetOrN>::from_iter_u8(data.query.sequences.query.iter().copied())
.unwrap();
template_switch_distance_a_star_align::<
AlignmentStrategySelection<
DnaAlphabetOrN,
U64Cost,
AntiDiagonalNodeOrdStrategy,
ML,
NoChainingStrategy<U64Cost>,
TS,
AllowSecondaryDeletionStrategy,
NoShortcutStrategy<U64Cost>,
AllowPrimaryMatchStrategy,
NoPrunePrimaryRangeStrategy,
MaxTemplateSwitchTotalLengthStrategy,
DS,
>,
_,
>(
reference.as_genome_subsequence(),
query.as_genome_subsequence(),
"reference",
"query",
data.query.ranges,
AdditionalExplicitTSMStartsAndEnds::default(),
data.costs,
DynamicStrategies {
ts_14_out_of_range: Ts14OutOfRangeStrategy::default(),
},
None,
Some(data.memory),
false,
true,
count_mem,
)
}