setup.rs

use std::fs::File;
use std::io::Read;
use std::path::Path;
use std::str::FromStr;

use std::ops::{Index, IndexMut};

use nom::Parser;

use crate::rust_cfd;

pub const GAMMA: f32 = 1.4;
pub const ff_mach: f32 = 1.2;
pub const NDIM: usize = 3;
pub const NNB: usize = 4;
pub const RK: usize = 3;
pub const VAR_DENSITY: usize = 0;
pub const VAR_MOMENTUM: usize = 1;
pub const VAR_DENSITY_ENERGY: usize = VAR_MOMENTUM + NDIM;
pub const NVAR: usize = VAR_DENSITY_ENERGY + 1;
pub const deg_angle_of_attack: f32 = 0.0;

#[repr(align(64))]
struct Alignment([u8; 64]);

// An aligned slice is stored as a boxed slice and an offset number of elements
// that we skip over to get the desired alignment (of 64 bytes)
pub struct AlignedSlice<T> {
    slice: Box<[T]>,
    offset: usize,
    length: usize,
}

impl<T> AlignedSlice<T>
where
    T: Default,
{
    pub fn of_len(len: usize) -> Self {
        // The maximum number of elements that may be waisted in getting the alignment we need is
        // (64 - alignment of T) / size of T
        let extra_elements = (64 - std::mem::align_of::<T>()) / std::mem::size_of::<T>();
        let slice: Box<[T]> = (0..len + extra_elements)
            .map(|_| Default::default())
            .collect();
        let offset = unsafe { slice.align_to::<Alignment>().0.len() };
        assert!(offset <= extra_elements);

        AlignedSlice {
            slice,
            offset,
            length: len,
        }
    }
}

impl<T> AlignedSlice<T> {
    pub fn as_slice(&self) -> &[T] {
        &self.slice[self.offset..self.offset + self.length]
    }

    pub fn as_mut_slice(&mut self) -> &mut [T] {
        &mut self.slice[self.offset..self.offset + self.length]
    }
}

impl<T> Index<usize> for AlignedSlice<T> {
    type Output = T;

    fn index(&self, index: usize) -> &Self::Output {
        &self.slice[index + self.offset]
    }
}

impl<T> IndexMut<usize> for AlignedSlice<T> {
    fn index_mut(&mut self, index: usize) -> &mut Self::Output {
        &mut self.slice[index + self.offset]
    }
}

// TODO: I don't think Default should be needed here, but it would require some
// other constructor probably one that uses Clone an an element rather than Default
impl<T> Clone for AlignedSlice<T>
where
    T: Clone + Default,
{
    fn clone(&self) -> Self {
        let mut new_slice = Self::of_len(self.length);
        new_slice.as_mut_slice().clone_from_slice(self.as_slice());
        new_slice
    }
}

#[repr(C)]
pub struct Float3 {
    pub x: f32,
    pub y: f32,
    pub z: f32,
}

pub struct FarFieldConditions {
    pub ff_variable: AlignedSlice<f32>,
    pub ff_fc_momentum_x: Float3,
    pub ff_fc_momentum_y: Float3,
    pub ff_fc_momentum_z: Float3,
    pub ff_fc_density_energy: Float3,
}

pub fn set_far_field_conditions() -> FarFieldConditions {
    let mut ff_variable: AlignedSlice<f32> = AlignedSlice::of_len(NVAR);

    let angle_of_attack = std::f32::consts::PI / 180.0 * deg_angle_of_attack;

    ff_variable[VAR_DENSITY] = 1.4;

    let ff_pressure = 1.0;
    let ff_speed_of_sound = (GAMMA * ff_pressure / ff_variable[VAR_DENSITY]).sqrt();
    let ff_speed = ff_mach * ff_speed_of_sound;

    let ff_velocity = Float3 {
        x: ff_speed * angle_of_attack.cos(),
        y: ff_speed * angle_of_attack.sin(),
        z: 0.0,
    };

    ff_variable[VAR_MOMENTUM + 0] = ff_variable[VAR_DENSITY] * ff_velocity.x;
    ff_variable[VAR_MOMENTUM + 1] = ff_variable[VAR_DENSITY] * ff_velocity.y;
    ff_variable[VAR_MOMENTUM + 2] = ff_variable[VAR_DENSITY] * ff_velocity.z;

    ff_variable[VAR_DENSITY_ENERGY] =
        ff_variable[VAR_DENSITY] * (0.5 * (ff_speed * ff_speed)) + (ff_pressure / (GAMMA - 1.0));

    let ff_momentum = Float3 {
        x: ff_variable[VAR_MOMENTUM + 0],
        y: ff_variable[VAR_MOMENTUM + 1],
        z: ff_variable[VAR_MOMENTUM + 2],
    };
    let (ff_fc_momentum_x, ff_fc_momentum_y, ff_fc_momentum_z, ff_fc_density_energy) =
        rust_cfd::compute_flux_contribution(
            ff_variable[VAR_DENSITY],
            &ff_momentum,
            ff_variable[VAR_DENSITY_ENERGY],
            ff_pressure,
            &ff_velocity,
        );

    FarFieldConditions {
        ff_variable,
        ff_fc_momentum_x,
        ff_fc_momentum_y,
        ff_fc_momentum_z,
        ff_fc_density_energy,
    }
}

pub struct GeometryData {
    pub nelr: usize,
    pub areas: AlignedSlice<f32>,
    pub elements_surrounding_elements: AlignedSlice<i32>,
    pub normals: AlignedSlice<f32>,
}

pub fn read_domain_geometry<P>(path: P, block_size: usize) -> GeometryData
where
    P: AsRef<Path>,
{
    let mut file = File::open(path).expect("Error opening input file");
    let mut contents = String::new();
    file.read_to_string(&mut contents)
        .expect("Error reading input file");

    let mut parser = nom::combinator::all_consuming(|s| cfd_parser(block_size, s));
    let (_, result) = parser.parse(&contents).expect("Parser error");

    result
}

fn cfd_parser<'a>(block_size: usize, text: &'a str) -> nom::IResult<&'a str, GeometryData> {
    let text = nom::character::complete::multispace0(text)?.0;
    let (text, nel) = nom::character::complete::digit1(text)?;

    let nel = usize::from_str(nel).unwrap();
    let nelr = block_size * ((nel / block_size) + usize::min(1, nel % block_size));

    let mut areas: AlignedSlice<f32> = AlignedSlice::of_len(nelr);
    let mut elements_surrounding_elements: AlignedSlice<i32> = AlignedSlice::of_len(nelr * NNB);
    let mut normals: AlignedSlice<f32> = AlignedSlice::of_len(NDIM * NNB * nelr);

    // read in data
    let mut text = text;
    for i in 0..nel {
        let t = nom::character::complete::multispace0(text)?.0;
        let (t, val) = nom::number::complete::float(t)?;
        text = t;

        areas[i] = val;

        for j in 0..NNB {
            let t = nom::character::complete::multispace0(text)?.0;
            let (t, neg) = if let Ok((t, _)) =
                nom::character::complete::char::<&str, nom::error::Error<&str>>('-')(t)
            {
                (t, true)
            } else {
                (t, false)
            };
            let (t, val) = nom::character::complete::digit1(t)?;
            text = t;

            let val = i32::from_str(val).unwrap();
            let val = if neg { -val } else { val };
            elements_surrounding_elements[i + j * nelr] = if val < 0 { -1 } else { val } - 1; // it's coming in with Fortran numbering

            for k in 0..NDIM {
                let t = nom::character::complete::multispace0(text)?.0;
                let (t, val) = nom::number::complete::float(t)?;
                text = t;

                normals[i + (j + k * NNB) * nelr] = -val;
            }
        }
    }

    // fill in remaining data
    let last = nel - 1;
    for i in nel..nelr {
        areas[i] = areas[last];
        for j in 0..NNB {
            elements_surrounding_elements[i + j * nelr] =
                elements_surrounding_elements[last + j * nelr];
            for k in 0..NDIM {
                normals[i + (j + k * NNB) * nelr] = normals[last + (j + k * NNB) * nelr];
            }
        }
    }

    let text = nom::character::complete::multispace0(text)?.0;
    Ok((
        text,
        GeometryData {
            nelr,
            areas,
            elements_surrounding_elements,
            normals,
        },
    ))
}

pub fn initialize_variables(nelr: usize, ff_variable: &[f32]) -> AlignedSlice<f32> {
    let mut variables = AlignedSlice::of_len(nelr * NVAR);

    for i in 0..nelr {
        for j in 0..NVAR {
            variables[i + j * nelr] = ff_variable[j];
        }
    }

    variables
}