#![allow(dead_code)]
use {
crate::{
Ident,
database::{
chunk::RecordWriter,
Partition,
PartitionKey,
storage::{
PartitionsBuilder,
RecordStorage,
},
},
hash::ContentHasher,
record::{
LaburnumRecord,
Record,
},
scheduler::{
lanes::DEFAULT_LANE,
tests::common::wait_for,
},
},
serde::Serialize,
std::hash::Hash,
};
#[derive(Debug, Clone, Hash, Serialize)]
pub struct CompilationData {
pub data: Vec<u64>,
}
#[derive(Debug, Clone, Copy, Hash, PartialEq, Eq, Serialize)]
pub struct CompilationIdx(usize);
impl From<LaburnumRecord> for CompilationData {
fn from(_record: LaburnumRecord) -> Self {
panic!("CompilationData does not support LaburnumRecord conversion")
}
}
impl Record for CompilationData {
fn content_hash(&self) -> crate::ContentHash {
crate::record::hash_record(&self.data)
}
}
impl bluegum::Bluegum for CompilationData {
fn node(&self, b: &mut bluegum::Builder) {
b.name("CompilationData")
.field("data", format!("{:?}", self.data));
}
}
impl bluegum::BluegumWithState<dyn crate::SpanResolver> for CompilationData {}
impl<P: crate::database::storage::Partitions> crate::record::CollectReferences<P>
for CompilationData
{
}
#[derive(Debug, Clone, serde::Serialize)]
pub struct CompilationRef<'a> {
pub data: &'a [u64],
}
impl<'a> crate::record::LaburnumRecordRef for CompilationRef<'a> {
fn as_laburnum_record(&self) -> Option<&LaburnumRecord> {
None
}
fn serialize_with_source_cache<P, T, Ser>(
&self,
_source_cache: &crate::source::SourceCache<P, T>,
serializer: Ser,
) -> Result<Ser::Ok, Ser::Error>
where
P: crate::database::storage::Partitions,
T: crate::protocol::lsp::LanguageServer<P>,
Ser: serde::Serializer,
{
serde::Serialize::serialize(self, serializer)
}
}
#[derive(Debug, Default, Clone)]
pub struct CompilationStorage {
records: Vec<Vec<u64>>,
}
impl RecordStorage for CompilationStorage {
type Builder = CompilationStorageBuilder;
type Index = CompilationIdx;
type RecordRef<'a> = CompilationRef<'a>;
fn get(&self, idx: &Self::Index) -> Option<Self::RecordRef<'_>> {
self.records.get(idx.0).map(|data| {
CompilationRef {
data: data.as_slice(),
}
})
}
fn hash_contents(&self, hasher: &mut ContentHasher) {
for record in &self.records {
for val in record {
hasher.update(&val.to_le_bytes());
}
}
}
}
#[derive(Debug, Default)]
pub struct CompilationStorageBuilder {
records: Vec<Vec<u64>>,
}
impl PartitionsBuilder for CompilationStorageBuilder {
type Record = CompilationData;
type Storage = CompilationStorage;
fn push(&mut self, record: Self::Record) -> CompilationIdx {
let idx = self.records.len();
self.records.push(record.data);
CompilationIdx(idx)
}
fn build(self) -> Self::Storage {
CompilationStorage {
records: self.records,
}
}
}
pub struct CompilationPartition;
impl PartitionKey for CompilationPartition {
const KEY: Ident = Ident::new("compilation_dag_test::records");
}
impl Partition for CompilationPartition {
type Record = CompilationData;
type IndexEntry = crate::database::partitions::HandleEntry<Self>;
type SortKey = String;
fn index_entry_from_handle(
handle: crate::database::RecordHandle<Self>,
) -> Self::IndexEntry {
crate::database::partitions::HandleEntry::new(handle)
}
}
crate::define_partitions! {
Compilation,
partitions = [Compilation,],
}
fn test_scheduler() -> (
std::sync::Arc<
crate::scheduler::Scheduler<
CompilationPartitions,
crate::server::LaburnumLanguageServer,
>,
>,
crate::connect::ipc::Connection,
) {
let (server_conn, client_conn) = crate::connect::ipc::Connection::memory();
let filesystems = std::sync::Arc::new(parking_lot::RwLock::new(Vec::new()));
let source_cache =
std::sync::Arc::new(parking_lot::RwLock::new(crate::SourceCache::new()));
let scheduler = crate::scheduler::Scheduler::new_with_worker_count(
server_conn,
std::sync::Arc::new(crate::server::LaburnumLanguageServer),
filesystems,
source_cache,
1,
);
(scheduler, client_conn)
}
#[macro_rules_attribute::apply(smol_macros::test!)]
#[test_log::test]
async fn test_simple_linear_compilation() {
let (scheduler, _conn) = test_scheduler();
let lex_name = "file1_lex";
let parse_name = "file1_parse";
let typecheck_name = "file1_typecheck";
let ir_name = "file1_ir";
let codegen_name = "file1_codegen";
let _lex_id = Ident::new(lex_name);
let _parse_id = Ident::new(parse_name);
let _typecheck_id = Ident::new(typecheck_name);
let _ir_id = Ident::new(ir_name);
scheduler.queue(
move |_ctx| {
let task_id = Ident::new(lex_name);
async move {
let mut writer = RecordWriter::new(task_id);
writer.insert::<CompilationPartition, _>(lex_name.to_string(), CompilationData {
data: vec![1],
});
Some(writer)
}
},
DEFAULT_LANE,
);
scheduler.queue(
move |mut ctx| {
let task_id = Ident::new(parse_name);
async move {
let query_client = ctx.query_client();
let results =
query_client.get_record(CompilationPartition::KEY, lex_name.to_string()).await;
let mut combined_data = Vec::new();
if !results.is_empty()
&& let Some(CompilationRecord::Compilation(record)) = results.get(&results.records()[0])
{
combined_data.extend(record.data);
}
let mut writer = RecordWriter::new(task_id);
writer.insert::<CompilationPartition, _>(parse_name.to_string(), CompilationData {
data: combined_data,
});
Some(writer)
}
},
DEFAULT_LANE,
);
scheduler.queue(
move |mut ctx| {
let task_id = Ident::new(typecheck_name);
async move {
let query_client = ctx.query_client();
let results = query_client
.get_record(CompilationPartition::KEY, parse_name.to_string())
.await;
let mut combined_data = Vec::new();
if !results.is_empty()
&& let Some(CompilationRecord::Compilation(record)) = results.get(&results.records()[0])
{
combined_data.extend(record.data);
}
let mut writer = RecordWriter::new(task_id);
writer.insert::<CompilationPartition, _>(typecheck_name.to_string(), CompilationData {
data: combined_data,
});
Some(writer)
}
},
DEFAULT_LANE,
);
scheduler.queue(
move |mut ctx| {
let task_id = Ident::new(ir_name);
async move {
let query_client = ctx.query_client();
let results = query_client
.get_record(CompilationPartition::KEY, typecheck_name.to_string())
.await;
let mut combined_data = Vec::new();
if !results.is_empty()
&& let Some(CompilationRecord::Compilation(record)) = results.get(&results.records()[0])
{
combined_data.extend(record.data);
}
let mut writer = RecordWriter::new(task_id);
writer.insert::<CompilationPartition, _>(ir_name.to_string(), CompilationData {
data: combined_data,
});
Some(writer)
}
},
DEFAULT_LANE,
);
scheduler.queue(
move |mut ctx| {
let task_id = Ident::new(codegen_name);
async move {
let query_client = ctx.query_client();
let results = query_client.get_record(CompilationPartition::KEY, ir_name.to_string()).await;
let mut combined_data = Vec::new();
if !results.is_empty()
&& let Some(CompilationRecord::Compilation(record)) = results.get(&results.records()[0])
{
combined_data.extend(record.data);
}
let mut writer = RecordWriter::new(task_id);
writer.insert::<CompilationPartition, _>(codegen_name.to_string(), CompilationData {
data: combined_data,
});
Some(writer)
}
},
DEFAULT_LANE,
);
scheduler.spawn_workers();
let verify_success =
std::sync::Arc::new(std::sync::atomic::AtomicBool::new(false));
let verify_ok = verify_success.clone();
scheduler.queue(
move |mut ctx| {
let success = verify_ok.clone();
async move {
let query_client = ctx.query_client();
let _codegen_id = Ident::new(codegen_name);
let codegen_records = query_client
.get_record(CompilationPartition::KEY, codegen_name.to_string())
.await;
if !codegen_records.is_empty() && codegen_records.len() == 1 {
success.store(true, std::sync::atomic::Ordering::SeqCst);
}
None
}
},
DEFAULT_LANE,
);
assert!(
wait_for(
|| verify_success.load(std::sync::atomic::Ordering::SeqCst),
std::time::Duration::from_secs(5)
)
.await,
"Codegen task should have completed"
);
}
#[macro_rules_attribute::apply(smol_macros::test!)]
#[test_log::test]
async fn test_two_file_compilation_with_import() {
let (scheduler, _conn) = test_scheduler();
let file1_lex = "file1_lex";
let file1_parse = "file1_parse";
let file1_typecheck = "file1_typecheck";
let file2_lex = "file2_lex";
let file2_parse = "file2_parse";
let file2_typecheck = "file2_typecheck";
let _file1_lex_id = Ident::new(file1_lex);
let _file1_parse_id = Ident::new(file1_parse);
let _file1_typecheck_id = Ident::new(file1_typecheck);
let _file2_lex_id = Ident::new(file2_lex);
let _file2_parse_id = Ident::new(file2_parse);
scheduler.queue(
move |_ctx| {
let task_id = Ident::new(file1_lex);
async move {
let mut writer = RecordWriter::new(task_id);
writer.insert::<CompilationPartition, _>(file1_lex.to_string(), CompilationData {
data: vec![1],
});
Some(writer)
}
},
DEFAULT_LANE,
);
scheduler.queue(
move |mut ctx| {
let task_id = Ident::new(file1_parse);
async move {
let query_client = ctx.query_client();
let results = query_client
.get_record(CompilationPartition::KEY, file1_lex.to_string())
.await;
let mut combined_data = Vec::new();
if !results.is_empty()
&& let Some(CompilationRecord::Compilation(record)) = results.get(&results.records()[0])
{
combined_data.extend(record.data);
}
let mut writer = RecordWriter::new(task_id);
writer.insert::<CompilationPartition, _>(file1_parse.to_string(), CompilationData {
data: combined_data,
});
Some(writer)
}
},
DEFAULT_LANE,
);
scheduler.queue(
move |mut ctx| {
let task_id = Ident::new(file1_typecheck);
async move {
let query_client = ctx.query_client();
let results = query_client
.get_record(CompilationPartition::KEY, file1_parse.to_string())
.await;
let mut combined_data = Vec::new();
if !results.is_empty()
&& let Some(CompilationRecord::Compilation(record)) = results.get(&results.records()[0])
{
combined_data.extend(record.data);
}
let mut writer = RecordWriter::new(task_id);
writer.insert::<CompilationPartition, _>(file1_typecheck.to_string(), CompilationData {
data: combined_data,
});
Some(writer)
}
},
DEFAULT_LANE,
);
scheduler.queue(
move |_ctx| {
let task_id = Ident::new(file2_lex);
async move {
let mut writer = RecordWriter::new(task_id);
writer.insert::<CompilationPartition, _>(file2_lex.to_string(), CompilationData {
data: vec![2],
});
Some(writer)
}
},
DEFAULT_LANE,
);
scheduler.queue(
move |mut ctx| {
let task_id = Ident::new(file2_parse);
async move {
let query_client = ctx.query_client();
let results = query_client
.get_record(CompilationPartition::KEY, file2_lex.to_string())
.await;
let mut combined_data = Vec::new();
if !results.is_empty()
&& let Some(CompilationRecord::Compilation(record)) = results.get(&results.records()[0])
{
combined_data.extend(record.data);
}
let mut writer = RecordWriter::new(task_id);
writer.insert::<CompilationPartition, _>(file2_parse.to_string(), CompilationData {
data: combined_data,
});
Some(writer)
}
},
DEFAULT_LANE,
);
scheduler.queue(
move |mut ctx| {
let task_id = Ident::new(file2_typecheck);
async move {
let query_client = ctx.query_client();
let parse_results = query_client
.get_record(CompilationPartition::KEY, file2_parse.to_string())
.await;
let file1_results = query_client
.get_record(CompilationPartition::KEY, file1_typecheck.to_string())
.await;
let mut combined_data = Vec::new();
if !parse_results.is_empty()
&& let Some(CompilationRecord::Compilation(record)) = parse_results.get(&parse_results.records()[0])
{
combined_data.extend(record.data);
}
if !file1_results.is_empty()
&& let Some(CompilationRecord::Compilation(record)) = file1_results.get(&file1_results.records()[0])
{
combined_data.extend(record.data);
}
let mut writer = RecordWriter::new(task_id);
writer.insert::<CompilationPartition, _>(file2_typecheck.to_string(), CompilationData {
data: combined_data,
});
Some(writer)
}
},
DEFAULT_LANE,
);
scheduler.spawn_workers();
let verify_success =
std::sync::Arc::new(std::sync::atomic::AtomicBool::new(false));
let verify_ok = verify_success.clone();
scheduler.queue(
move |mut ctx| {
let success = verify_ok.clone();
async move {
let query_client = ctx.query_client();
let _file1_typecheck_id = Ident::new(file1_typecheck);
let _file2_typecheck_id = Ident::new(file2_typecheck);
let file1_records = query_client
.get_record(CompilationPartition::KEY, file1_typecheck.to_string())
.await;
let file2_records = query_client
.get_record(CompilationPartition::KEY, file2_typecheck.to_string())
.await;
if !file1_records.is_empty()
&& file1_records.len() == 1
&& !file2_records.is_empty()
&& file2_records.len() == 1
{
success.store(true, std::sync::atomic::Ordering::SeqCst);
}
None
}
},
DEFAULT_LANE,
);
assert!(
wait_for(
|| verify_success.load(std::sync::atomic::Ordering::SeqCst),
std::time::Duration::from_secs(5)
)
.await,
"File1 and File2 typecheck should have completed"
);
}