shaderc/src/compiler/backend/mod.rs

#[cfg(test)]
mod tests;

use crate::compiler::{Instruction, Module};
use std::fmt;
use std::fmt::Write;

use super::StorageClass;

use crate::parser::Ast;

#[derive(Debug, PartialEq, Clone)]
pub enum Type {
    Pointer(Box<Type>, StorageClass),
    Vector(Box<Type>, u32),
    Float(u32),
    Int(u32),
    Unsigned(u32),
    Void,
}

impl fmt::Display for Type {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        match self {
            Type::Void => write!(f, "<>"),
            Type::Pointer(typ, storage_class) => match storage_class {
                StorageClass::Input => write!(f, "*{}i", typ),
                StorageClass::Output => write!(f, "*{}o", typ),
                StorageClass::Undefined => panic!("Bound a non-declared variable"),
            },
            Type::Vector(typ, size) => write!(f, "v{}{}", size, typ),
            Type::Float(size) => write!(f, "f{}", size),
            Type::Int(size) => write!(f, "s{}", size),
            Type::Unsigned(size) => write!(f, "u{}", size),
        }
    }
}

pub fn parse_type(typ: &String) -> Type {
    // pointers have the format *<t>i or *<t>o, for the storage class
    // floats have the format f<size>, so f32, f64...
    // ints have the format s<size>, so s32, s64...
    // unsigned ints have the format u<size>, so u32, u64...
    // So, *v4f32i is a pointer to a vector of 4 floats, with the storage class Input.
    // *v4f32o is a pointer to a vector of 4 floats, with the storage class Output.

    let c = typ.chars().next().unwrap();
    match c {
        '<' => {
            assert_eq!(typ.as_str(), "<>");
            Type::Void
        }
        '*' => {
            let mut chars = typ.chars();
            chars.next();
            let typ = chars.collect::<String>();
            let storage_class = match typ.chars().last().unwrap() {
                'i' => StorageClass::Input,
                'o' => StorageClass::Output,
                _ => panic!("Invalid storage class"),
            };
            let typ = typ.chars().take(typ.len() - 1).collect::<String>();
            Type::Pointer(Box::new(parse_type(&typ)), storage_class)
        }
        'v' => {
            let mut chars = typ.chars();
            chars.next();
            let size = chars.next().unwrap().to_digit(10).unwrap();
            let typ = chars.collect::<String>();
            Type::Vector(Box::new(parse_type(&typ)), size)
        }
        'f' => {
            let size = typ.chars().skip(1).collect::<String>().parse().unwrap();
            Type::Float(size)
        }
        's' => {
            let size = typ.chars().skip(1).collect::<String>().parse().unwrap();
            Type::Int(size)
        }
        'u' => {
            let size = typ.chars().skip(1).collect::<String>().parse().unwrap();
            Type::Unsigned(size)
        }
        _ => panic!("Invalid type"),
    }
}

fn emit_type(typ: Type, ops: &mut Vec<(String, String)>) {
    match &typ {
        Type::Void => {
            ops.push(("%void".to_string(), "OpTypeVoid".to_string()));
        }
        Type::Unsigned(size) => {
            ops.push((
                format!("%u{}", size).to_string(),
                format!("OpTypeInt {}", size),
            ));
        }
        Type::Int(size) => {
            ops.push((
                format!("%s{}", size).to_string(),
                format!("OpTypeInt {}", size),
            ));
        }
        Type::Float(size) => {
            ops.push((
                format!("%f{}", size).to_string(),
                format!("OpTypeFloat {}", size),
            ));
        }
        Type::Vector(in_typ, size) => {
            emit_type(*in_typ.clone(), ops);
            ops.push((
                fix_name(&typ.to_string()),
                format!(
                    "OpTypeVector {} {}",
                    fix_name(&in_typ.clone().to_string()),
                    size
                ),
            ))
        }
        Type::Pointer(in_typ, storage_class) => {
            emit_type(*in_typ.clone(), ops);
            let typ_id = fix_name(&typ.to_string());
            let storage_class = match storage_class {
                StorageClass::Input => "Input",
                StorageClass::Output => "Output",
                StorageClass::Undefined => panic!("Bound a non-declared variable"),
            };
            ops.push((
                typ_id,
                format!(
                    "OpTypePointer {} {}",
                    storage_class,
                    fix_name(&in_typ.to_string())
                ),
            ))
        }
    }
}

fn fix_name(name: &String) -> String {
    format!(
        "%{}",
        name.clone()
            .replace("-", "_")
            .replace("*", "p")
            .replace("<>", "void")
    )
}

fn has_id(name: String, ops: &Vec<(Option<String>, Vec<String>)>) -> bool {
    for op in ops {
        if op.0.is_some() && op.0.clone().unwrap() == name {
            return true;
        }
    }
    false
}

pub fn spirv_meta(module: Module) -> String {
    let mut spirv_asm = String::new();
    let mut ops: Vec<(Option<String>, Vec<String>)> = Vec::new();

    let capabilities: Vec<String> = module
        .capabilities
        .iter()
        .map(|c| format!("{:?}", c))
        .collect();
    for cap in capabilities {
        ops.push((None, vec!["OpCapability".to_string(), cap]));
    }

    let memory_model_address = match module.memory_model.addressing_model {
        crate::compiler::AddressingModel::Logical => "Logical",
        crate::compiler::AddressingModel::Physical32 => "Physical32",
        crate::compiler::AddressingModel::Physical64 => "Physical64",
        crate::compiler::AddressingModel::PhysicalStorageBuffer64 => "PhysicalStorageBuffer64",
    };
    let memory_model_model = match module.memory_model.memory_model {
        crate::compiler::MemoryModel::Simple => "Simple",
        crate::compiler::MemoryModel::GLSL450 => "GLSL450",
        crate::compiler::MemoryModel::OpenCL => "OpenCL",
        _ => todo!(),
    };
    ops.push((
        None,
        vec![
            "OpMemoryModel".to_string(),
            memory_model_address.to_string(),
            memory_model_model.to_string(),
        ],
    ));

    for entry in module.entry_points {
        let exec_model = match entry.execution_model {
            crate::compiler::ExecutionModel::Fragment => "Fragment",
            crate::compiler::ExecutionModel::Vertex => "Vertex",
        };
        let name = entry.name;
        let interface: Vec<String> = entry
            .interface
            .iter()
            .map(|i| fix_name(&i.to_string()))
            .collect();
        let exec_mode = match entry.execution_mode {
            crate::compiler::ExecutionMode::OriginUpperLeft => "OriginUpperLeft",
        };
        ops.push((
            None,
            vec![
                "OpEntryPoint".to_string(),
                exec_model.to_string(),
                fix_name(&name),
                format!("\"{}\"", name),
                interface.join(" "),
            ],
        ));
        ops.push((
            None,
            vec![
                "OpExecutionMode".to_string(),
                fix_name(&name),
                exec_mode.to_string(),
            ],
        ));
    }

    for global in module.globals {
        let name = fix_name(&global.name);
        let _typ = global.typ;
        let storage_class = match global.storage_class {
            crate::compiler::StorageClass::Input => "Input",
            crate::compiler::StorageClass::Output => "Output",
            crate::compiler::StorageClass::Undefined => panic!("Bound a non-declared variable"),
        };
        let mut type_ops = Vec::new();
        emit_type(parse_type(&_typ), &mut type_ops);
        for op in type_ops {
            if has_id(op.0.clone(), &ops) {
                continue;
            }
            ops.push((Some(op.0), vec![op.1]));
        }
        ops.push((
            Some(name.clone()),
            vec![
                "OpVariable".to_string(),
                fix_name(&_typ),
                storage_class.to_string(),
            ],
        ));
        for dec in global.decorations {
            // Decorations have the format Location 0, or Builtin FragCoord
            let dec = match dec {
                crate::compiler::Decoration::Location(loc) => format!("Location {}", loc),
                crate::compiler::Decoration::BuiltIn(builtin) => {
                    let builtin = match builtin {
                        crate::compiler::BuiltinDecoration::FragCoord => "FragCoord",
                    };
                    format!("BuiltIn {}", builtin)
                }
            };
            ops.push((None, vec!["OpDecorate".to_string(), name.clone(), dec]));
        }
    }

    for fun in module.functions {
        let name = fix_name(&fun.name);
        let return_type = fix_name(&fun.return_type);
        let mut type_ops = Vec::new();
        emit_type(parse_type(&fun.return_type), &mut type_ops);
        for op in type_ops {
            if has_id(op.0.clone(), &ops) {
                continue;
            }
            ops.push((Some(op.0), vec![op.1]));
        }
        // Push OpFunctionType
        ops.push((
            Some(name.clone()),
            vec!["OpTypeFunction".to_string(), return_type.clone()],
        ));
    }

    for op in ops {
        if op.0.is_some() {
            write!(spirv_asm, "{} = ", op.0.unwrap()).unwrap();
        }
        for arg in op.1 {
            write!(spirv_asm, "{} ", arg).unwrap();
        }
        writeln!(spirv_asm).unwrap();
    }
    spirv_asm
}

enum Number {
    Int(i32),
    Float(f32),
    NotANumber,
}

fn match_number(s: &str) -> Number {
    // floats have to be in the format of [0-9]+\.[0-9]*
    // integers have to be in the format of [0-9]+
    let mut chars = s.chars();
    let mut has_dot = false;
    while let Some(c) = chars.next() {
        if c == '.' {
            if has_dot {
                // only one dot allowed.
                return Number::NotANumber;
            }
            has_dot = true;
        } else if !c.is_digit(10) {
            // not a number;
            // has a character that is not a digit or a dot.
            return Number::NotANumber;
        }
    }
    if has_dot {
        // we have checked that the whole thing is numbers and dots
        // and that there is precisely one dot
        // that matches the regex.
        Number::Float(s.parse().unwrap())
    } else {
        // this is an integer, since no dots.
        Number::Int(s.parse().unwrap())
    }
}

fn compile_biop(
    op: &str,
    lst: &mut Vec<Ast>,
    vars: &mut Vec<(String, String)>,
    constants: &mut Vec<(String, String)>,
    types: &mut Vec<String>,
    counter: &mut i32,
    stack: &mut Vec<String>,
    out: &mut Vec<(Option<String>, String)>,
) {
    assert!(lst.len() == 2);
    let rhs = lst.pop().unwrap();
    let lhs = lst.pop().unwrap();
    compile_ast_ssa(lhs, vars, constants, types, counter, stack, out);
    compile_ast_ssa(rhs, vars, constants, types, counter, stack, out);
    let rhs_id = stack.pop().unwrap();
    let lhs_id = stack.pop().unwrap();
    let id = String::from(counter.to_string());
    *counter += 1;
    out.push((
        Some(id.clone()),
        format!(
            "{} {} {} {}",
            op,
            fix_name(&String::from("f32")),
            fix_name(&lhs_id),
            fix_name(&rhs_id),
        ),
    ));
    stack.push(id);
}

pub fn compile_ast_ssa(
    ast: Ast,
    vars: &mut Vec<(String, String)>,
    constants: &mut Vec<(String, String)>,
    types: &mut Vec<String>,
    counter: &mut i32,
    stack: &mut Vec<String>,
    out: &mut Vec<(Option<String>, String)>,
) {
    match ast.clone().list() {
        Some(l) => {
            let mut lst = l.clone();
            assert!(!lst.is_empty());
            let fun = lst.remove(0);
            assert!(true); // no safe remove, thanks Rust
            let fun_name = fun.symbol();
            assert!(fun_name.is_some());
            let fun_name = fun_name.unwrap();
            match fun_name.as_str() {
                "store-ptr" => {
                    assert!(lst.len() == 2);
                    let ptr = lst.pop().unwrap();
                    let val = lst.pop().unwrap();
                    compile_ast_ssa(ptr, vars, constants, types, counter, stack, out);
                    compile_ast_ssa(val, vars, constants, types, counter, stack, out);
                    let val_id = stack.pop().unwrap();
                    let ptr_id = stack.pop().unwrap();
                    out.push((
                        None,
                        format!("OpStore {} {}", fix_name(&val_id), fix_name(&ptr_id)),
                    ));
                }
                "/" => {
                    compile_biop(
                        "OpFDiv", &mut lst, vars, constants, types, counter, stack, out,
                    );
                }
                "*" => {
                    compile_biop(
                        "OpFMul", &mut lst, vars, constants, types, counter, stack, out,
                    );
                }
                "+" => {
                    compile_biop(
                        "OpFAdd", &mut lst, vars, constants, types, counter, stack, out,
                    );
                }
                "-" => {
                    compile_biop(
                        "OpFSub", &mut lst, vars, constants, types, counter, stack, out,
                    );
                }
                s => {
                    panic!(
                        "Unknown function: {} with params {:#?} in context:\n{:#?}",
                        s, lst, ast
                    );
                }
            }
        }
        None => {
            let sym = ast.clone().symbol();
            assert!(sym.is_some());
            let sym = sym.unwrap();
            match match_number(&sym) {
                Number::Int(i) => {
                    let key = format!("i32_{}", i);
                    let mut contains = false;
                    for c in constants.iter() {
                        if c.0 == key {
                            contains = true;
                        }
                    }
                    if !contains {
                        constants.push((key.clone(), format!("OpConstant %i32 {}", i.to_string())));
                    }
                    stack.push(key);
                }
                Number::Float(f) => {
                    let key = format!("f32_{}", f);
                    let mut contains = false;
                    for c in constants.iter() {
                        if c.0 == key {
                            contains = true;
                        }
                    }
                    if !contains {
                        constants.push((key.clone(), format!("OpConstant %f32 {}", f.to_string())));
                    }
                    for t in types.iter() {
                        if t == "f32" {
                            contains = true;
                        }
                    }
                    if !contains {
                        types.push("f32".to_string());
                    }
                    stack.push(key);
                }
                Number::NotANumber => {
                    for v in vars.iter() {
                        if v.0 == sym {
                            stack.push(v.0.clone());
                            return;
                        }
                    }
                    panic!("Unknown variable or constant: {}", sym);
                }
            }
        }
    }
}

pub fn compile_fun_ssa(module: &mut Module) {
    for fun in module.functions.iter_mut() {
        assert!(fun.ast.is_some());
        let ast = fun.ast.as_mut().unwrap();
        let block = ast.clone().list().unwrap().get(0).unwrap().clone();
        let mut vars = vec![];
        let mut constants = vec![];
        let mut types = vec![];
        let mut counter = Box::new(0);
        let mut stack = vec![];
        let mut out_op = vec![];
        for v in &module.globals {
            vars.push((v.name.clone(), v.typ.clone()));
        }
        compile_ast_ssa(
            block,
            &mut vars,
            &mut constants,
            &mut types,
            &mut counter,
            &mut stack,
            &mut out_op,
        );
        let mut out_pre = vec![];
        for t in &types {
            let typ = parse_type(t);
            let mut type_ops = vec![];
            emit_type(typ, &mut type_ops);
            for type_op in type_ops {
                out_pre.push((Some(type_op.0), type_op.1));
            }
        }
        for c in &constants {
            out_pre.push((Some(fix_name(&c.0.clone())), c.1.clone()));
        }
        let mut out_ops = out_pre.clone();
        for op in out_op {
            if op.0.is_some() {
                out_ops.push((Some(fix_name(&op.0.unwrap())), op.1));
            } else {
                out_ops.push((None, op.1));
            }
        }
        for op in out_ops {
            let split: Vec<String> = op.1.split(" ").map(|s| s.to_string()).collect();
            let op_name: String = (&split[0]).clone();
            let op_args: Vec<String> = split[1..].iter().map(|s| s.clone()).collect();
            let op_id = op.0.clone();
            let ins: Instruction = Instruction {
                result_id: op_id,
                op: op_name,
                operands: op_args,
            };
            fun.body.as_mut().unwrap().push(ins);
        }
    }
}