use std::sync::{Arc, OnceLock};
use serde::{Deserialize, Serialize};
use super::{DictionaryId, DictionaryResult, NoriDictionary, UserDictionary};
mod hash;
use crate::morphology::resources::{Artifact, Request, Resources};
pub use crate::morphology::resources::{DictionaryBytes, ResourceCacheStats, ResourceLimits};
pub use hash::ResourceHash;
pub const DEFAULT_NORI_DICTIONARY: &str = "lucene-10.5.1";
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
#[serde(rename_all = "snake_case")]
pub enum DictionaryRequest {
Name(String),
Sha256(ResourceHash),
}
impl std::fmt::Display for DictionaryRequest {
fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
match self {
Self::Name(name) => formatter.write_str(name),
Self::Sha256(hash) => write!(formatter, "sha256:{hash}"),
}
}
}
pub struct DictionaryArtifact {
pub sha256: ResourceHash,
pub bytes: DictionaryBytes,
}
pub trait DictionaryResolver: Send + Sync {
fn resolve(&self, request: &DictionaryRequest) -> DictionaryResult<Option<DictionaryArtifact>>;
}
impl<F> DictionaryResolver for F
where
F: Fn(&DictionaryRequest) -> DictionaryResult<Option<DictionaryArtifact>> + Send + Sync,
{
fn resolve(&self, request: &DictionaryRequest) -> DictionaryResult<Option<DictionaryArtifact>> {
self(request)
}
}
pub struct ResolvedDictionary {
sha256: ResourceHash,
bytes: DictionaryBytes,
model: Arc<NoriDictionary>,
}
impl std::fmt::Debug for ResolvedDictionary {
fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
formatter
.debug_struct("ResolvedDictionary")
.field("sha256", &self.sha256)
.field("model", &self.model)
.finish_non_exhaustive()
}
}
impl ResolvedDictionary {
pub fn sha256(&self) -> ResourceHash {
self.sha256
}
pub fn model(&self) -> &Arc<NoriDictionary> {
&self.model
}
pub fn bytes(&self) -> &[u8] {
self.bytes.as_ref()
}
}
pub struct ResolvedUserDictionary {
sha256: ResourceHash,
model_id: DictionaryId,
dictionary: Option<Arc<UserDictionary>>,
empty_source: Option<Arc<str>>,
}
impl ResolvedUserDictionary {
pub fn sha256(&self) -> ResourceHash {
self.sha256
}
pub fn model_id(&self) -> DictionaryId {
self.model_id
}
pub fn dictionary(&self) -> Option<&Arc<UserDictionary>> {
self.dictionary.as_ref()
}
pub fn source(&self) -> &str {
match &self.dictionary {
Some(dictionary) => dictionary.source(),
None => self.empty_source.as_deref().expect("retained empty rules"),
}
}
}
struct Inner {
resolver: Arc<dyn DictionaryResolver>,
resources: Resources<ResolvedDictionary, ResolvedUserDictionary, DictionaryId>,
}
#[derive(Clone)]
pub struct NoriResources(Arc<Inner>);
impl Default for NoriResources {
fn default() -> Self {
static RESOURCES: OnceLock<NoriResources> = OnceLock::new();
RESOURCES
.get_or_init(|| {
Self::with_resolver(Arc::new(BundledResolver), ResourceLimits::default())
})
.clone()
}
}
impl NoriResources {
pub fn with_resolver(resolver: Arc<dyn DictionaryResolver>, limits: ResourceLimits) -> Self {
Self(Arc::new(Inner {
resolver,
resources: Resources::new(limits),
}))
}
pub fn limits(&self) -> ResourceLimits {
self.0.resources.limits()
}
pub fn cache_stats(&self) -> ResourceCacheStats {
self.0.resources.stats()
}
pub fn load_default(&self) -> DictionaryResult<Arc<ResolvedDictionary>> {
self.load(&DictionaryRequest::Name(DEFAULT_NORI_DICTIONARY.into()))
}
pub fn load(&self, request: &DictionaryRequest) -> DictionaryResult<Arc<ResolvedDictionary>> {
let shared_request = match request {
DictionaryRequest::Name(name) => Request::Name(name),
DictionaryRequest::Sha256(hash) => Request::Sha256(hash.into_bytes()),
};
self.0.resources.load(
shared_request,
|| {
self.0.resolver.resolve(request).map(|artifact| {
artifact.map(|artifact| Artifact {
sha256: artifact.sha256.into_bytes(),
bytes: artifact.bytes,
})
})
},
|artifact| {
let model =
NoriDictionary::from_bytes(artifact.bytes.as_ref(), self.limits().dictionary)?;
Ok(ResolvedDictionary {
sha256: ResourceHash::from_bytes(artifact.sha256),
bytes: artifact.bytes,
model,
})
},
)
}
pub fn compile_user(
&self,
source: &str,
model: &NoriDictionary,
) -> DictionaryResult<Arc<ResolvedUserDictionary>> {
self.0.resources.compile_user(source, model.id(), |hash| {
let dictionary = UserDictionary::compile(source, model, self.limits().user_dictionary)?;
let empty_source = dictionary.is_none().then(|| Arc::from(source));
Ok(ResolvedUserDictionary {
sha256: ResourceHash::from_bytes(hash),
model_id: model.id(),
dictionary,
empty_source,
})
})
}
}
struct BundledResolver;
impl DictionaryResolver for BundledResolver {
fn resolve(&self, request: &DictionaryRequest) -> DictionaryResult<Option<DictionaryArtifact>> {
let hash = uqa_nori_data::BUNDLE_SHA256.parse()?;
let matches = match request {
DictionaryRequest::Name(name) => name == DEFAULT_NORI_DICTIONARY,
DictionaryRequest::Sha256(expected) => *expected == hash,
};
Ok(matches.then_some(DictionaryArtifact {
sha256: hash,
bytes: DictionaryBytes::Static(uqa_nori_data::BUNDLE),
}))
}
}
#[cfg(test)]
mod tests;