//! Defines `ObjectBackend`.

use crate::traps::{ObjectTrapSink, ObjectTrapSite};
use cranelift_codegen::binemit::{
    Addend, CodeOffset, NullStackmapSink, NullTrapSink, Reloc, RelocSink,
};
use cranelift_codegen::entity::SecondaryMap;
use cranelift_codegen::isa::TargetIsa;
use cranelift_codegen::{self, binemit, ir};
use cranelift_module::{
    Backend, DataContext, DataDescription, DataId, FuncId, Init, Linkage, ModuleNamespace,
    ModuleResult,
};
use object::write::{
    Object, Relocation, SectionId, StandardSection, Symbol, SymbolId, SymbolSection,
};
use object::{RelocationEncoding, RelocationKind, SymbolFlags, SymbolKind, SymbolScope};
use std::collections::HashMap;
use target_lexicon::PointerWidth;

#[derive(Debug)]
/// Setting to enable collection of traps. Setting this to `Enabled` in
/// `ObjectBuilder` means that `ObjectProduct` will contains trap sites.
pub enum ObjectTrapCollection {
    /// `ObjectProduct::traps` will be empty
    Disabled,
    /// `ObjectProduct::traps` will contain trap sites
    Enabled,
}

/// A builder for `ObjectBackend`.
pub struct ObjectBuilder {
    isa: Box<dyn TargetIsa>,
    name: String,
    collect_traps: ObjectTrapCollection,
    libcall_names: Box<dyn Fn(ir::LibCall) -> String>,
    function_alignment: u64,
}

impl ObjectBuilder {
    /// Create a new `ObjectBuilder` using the given Cranelift target, that
    /// can be passed to [`Module::new`](cranelift_module::Module::new).
    ///
    /// `collect_traps` setting determines whether trap information is collected in the
    /// `ObjectProduct`.
    ///
    /// The `libcall_names` function provides a way to translate `cranelift_codegen`'s `ir::LibCall`
    /// enum to symbols. LibCalls are inserted in the IR as part of the legalization for certain
    /// floating point instructions, and for stack probes. If you don't know what to use for this
    /// argument, use `cranelift_module::default_libcall_names()`.
    pub fn new(
        isa: Box<dyn TargetIsa>,
        name: String,
        collect_traps: ObjectTrapCollection,
        libcall_names: Box<dyn Fn(ir::LibCall) -> String>,
    ) -> ModuleResult<Self> {
        Ok(Self {
            isa,
            name,
            collect_traps,
            libcall_names,
            function_alignment: 1,
        })
    }

    /// Set the alignment used for functions.
    pub fn function_alignment(&mut self, alignment: u64) -> &mut Self {
        self.function_alignment = alignment;
        self
    }
}

/// A `ObjectBackend` implements `Backend` and emits ".o" files using the `object` library.
///
/// See the `ObjectBuilder` for a convenient way to construct `ObjectBackend` instances.
pub struct ObjectBackend {
    isa: Box<dyn TargetIsa>,
    object: Object,
    functions: SecondaryMap<FuncId, Option<SymbolId>>,
    data_objects: SecondaryMap<DataId, Option<SymbolId>>,
    traps: SecondaryMap<FuncId, Vec<ObjectTrapSite>>,
    libcalls: HashMap<ir::LibCall, SymbolId>,
    libcall_names: Box<dyn Fn(ir::LibCall) -> String>,
    collect_traps: ObjectTrapCollection,
    function_alignment: u64,
}

impl Backend for ObjectBackend {
    type Builder = ObjectBuilder;

    type CompiledFunction = ObjectCompiledFunction;
    type CompiledData = ObjectCompiledData;

    // There's no need to return individual artifacts; we're writing them into
    // the output file instead.
    type FinalizedFunction = ();
    type FinalizedData = ();

    type Product = ObjectProduct;

    /// Create a new `ObjectBackend` using the given Cranelift target.
    fn new(builder: ObjectBuilder) -> Self {
        let triple = builder.isa.triple();
        let mut object = Object::new(triple.binary_format, triple.architecture);
        object.add_file_symbol(builder.name.as_bytes().to_vec());
        Self {
            isa: builder.isa,
            object,
            functions: SecondaryMap::new(),
            data_objects: SecondaryMap::new(),
            traps: SecondaryMap::new(),
            libcalls: HashMap::new(),
            libcall_names: builder.libcall_names,
            collect_traps: builder.collect_traps,
            function_alignment: builder.function_alignment,
        }
    }

    fn isa(&self) -> &dyn TargetIsa {
        &*self.isa
    }

    fn declare_function(&mut self, id: FuncId, name: &str, linkage: Linkage) {
        let (scope, weak) = translate_linkage(linkage);

        if let Some(function) = self.functions[id] {
            let symbol = self.object.symbol_mut(function);
            symbol.scope = scope;
            symbol.weak = weak;
        } else {
            let symbol_id = self.object.add_symbol(Symbol {
                name: name.as_bytes().to_vec(),
                value: 0,
                size: 0,
                kind: SymbolKind::Text,
                scope,
                weak,
                section: SymbolSection::Undefined,
                flags: SymbolFlags::None,
            });
            self.functions[id] = Some(symbol_id);
        }
    }

    fn declare_data(
        &mut self,
        id: DataId,
        name: &str,
        linkage: Linkage,
        _writable: bool,
        _align: Option<u8>,
    ) {
        let (scope, weak) = translate_linkage(linkage);

        if let Some(data) = self.data_objects[id] {
            let symbol = self.object.symbol_mut(data);
            symbol.scope = scope;
            symbol.weak = weak;
        } else {
            let symbol_id = self.object.add_symbol(Symbol {
                name: name.as_bytes().to_vec(),
                value: 0,
                size: 0,
                kind: SymbolKind::Data,
                scope,
                weak,
                section: SymbolSection::Undefined,
                flags: SymbolFlags::None,
            });
            self.data_objects[id] = Some(symbol_id);
        }
    }

    fn define_function(
        &mut self,
        func_id: FuncId,
        _name: &str,
        ctx: &cranelift_codegen::Context,
        _namespace: &ModuleNamespace<Self>,
        code_size: u32,
    ) -> ModuleResult<ObjectCompiledFunction> {
        let mut code: Vec<u8> = vec![0; code_size as usize];
        let mut reloc_sink = ObjectRelocSink::default();
        let mut trap_sink = ObjectTrapSink::default();
        let mut stackmap_sink = NullStackmapSink {};

        if let ObjectTrapCollection::Enabled = self.collect_traps {
            unsafe {
                ctx.emit_to_memory(
                    &*self.isa,
                    code.as_mut_ptr(),
                    &mut reloc_sink,
                    &mut trap_sink,
                    &mut stackmap_sink,
                )
            };
        } else {
            let mut trap_sink = NullTrapSink {};
            unsafe {
                ctx.emit_to_memory(
                    &*self.isa,
                    code.as_mut_ptr(),
                    &mut reloc_sink,
                    &mut trap_sink,
                    &mut stackmap_sink,
                )
            };
        }

        let symbol = self.functions[func_id].unwrap();
        let section = self.object.section_id(StandardSection::Text);
        let offset = self
            .object
            .add_symbol_data(symbol, section, &code, self.function_alignment);
        self.traps[func_id] = trap_sink.sites;
        Ok(ObjectCompiledFunction {
            offset,
            size: code_size,
            section,
            relocs: reloc_sink.relocs,
        })
    }

    fn define_data(
        &mut self,
        data_id: DataId,
        _name: &str,
        writable: bool,
        align: Option<u8>,
        data_ctx: &DataContext,
        _namespace: &ModuleNamespace<Self>,
    ) -> ModuleResult<ObjectCompiledData> {
        let &DataDescription {
            ref init,
            ref function_decls,
            ref data_decls,
            ref function_relocs,
            ref data_relocs,
        } = data_ctx.description();

        let size = init.size();
        let mut data = Vec::with_capacity(size);
        match *init {
            Init::Uninitialized => {
                panic!("data is not initialized yet");
            }
            Init::Zeros { .. } => {
                data.resize(size, 0);
            }
            Init::Bytes { ref contents } => {
                data.extend_from_slice(contents);
            }
        }

        let reloc_size = match self.isa.triple().pointer_width().unwrap() {
            PointerWidth::U16 => 16,
            PointerWidth::U32 => 32,
            PointerWidth::U64 => 64,
        };
        let mut relocs = Vec::new();
        for &(offset, id) in function_relocs {
            relocs.push(RelocRecord {
                offset,
                name: function_decls[id].clone(),
                kind: RelocationKind::Absolute,
                encoding: RelocationEncoding::Generic,
                size: reloc_size,
                addend: 0,
            });
        }
        for &(offset, id, addend) in data_relocs {
            relocs.push(RelocRecord {
                offset,
                name: data_decls[id].clone(),
                kind: RelocationKind::Absolute,
                encoding: RelocationEncoding::Generic,
                size: reloc_size,
                addend,
            });
        }

        let symbol = self.data_objects[data_id].unwrap();
        let section = self.object.section_id(if writable {
            StandardSection::Data
        } else if relocs.is_empty() {
            StandardSection::ReadOnlyData
        } else {
            StandardSection::ReadOnlyDataWithRel
        });
        let offset =
            self.object
                .add_symbol_data(symbol, section, &data, u64::from(align.unwrap_or(1)));
        Ok(ObjectCompiledData {
            offset,
            section,
            relocs,
        })
    }

    fn write_data_funcaddr(
        &mut self,
        _data: &mut ObjectCompiledData,
        _offset: usize,
        _what: ir::FuncRef,
    ) {
        unimplemented!()
    }

    fn write_data_dataaddr(
        &mut self,
        _data: &mut ObjectCompiledData,
        _offset: usize,
        _what: ir::GlobalValue,
        _usize: binemit::Addend,
    ) {
        unimplemented!()
    }

    fn finalize_function(
        &mut self,
        _id: FuncId,
        func: &ObjectCompiledFunction,
        namespace: &ModuleNamespace<Self>,
    ) {
        for &RelocRecord {
            offset,
            ref name,
            kind,
            encoding,
            size,
            addend,
        } in &func.relocs
        {
            let offset = func.offset + u64::from(offset);
            let symbol = self.get_symbol(namespace, name);
            self.object
                .add_relocation(
                    func.section,
                    Relocation {
                        offset,
                        size,
                        kind,
                        encoding,
                        symbol,
                        addend,
                    },
                )
                .unwrap();
        }
    }

    fn get_finalized_function(&self, _func: &ObjectCompiledFunction) {
        // Nothing to do.
    }

    fn finalize_data(
        &mut self,
        _id: DataId,
        data: &ObjectCompiledData,
        namespace: &ModuleNamespace<Self>,
    ) {
        for &RelocRecord {
            offset,
            ref name,
            kind,
            encoding,
            size,
            addend,
        } in &data.relocs
        {
            let offset = data.offset + u64::from(offset);
            let symbol = self.get_symbol(namespace, name);
            self.object
                .add_relocation(
                    data.section,
                    Relocation {
                        offset,
                        size,
                        kind,
                        encoding,
                        symbol,
                        addend,
                    },
                )
                .unwrap();
        }
    }

    fn get_finalized_data(&self, _data: &ObjectCompiledData) {
        // Nothing to do.
    }

    fn publish(&mut self) {
        // Nothing to do.
    }

    fn finish(self) -> ObjectProduct {
        ObjectProduct {
            object: self.object,
            functions: self.functions,
            data_objects: self.data_objects,
            traps: self.traps,
        }
    }
}

impl ObjectBackend {
    // This should only be called during finalization because it creates
    // symbols for missing libcalls.
    fn get_symbol(
        &mut self,
        namespace: &ModuleNamespace<Self>,
        name: &ir::ExternalName,
    ) -> SymbolId {
        match *name {
            ir::ExternalName::User { .. } => {
                if namespace.is_function(name) {
                    let id = namespace.get_function_id(name);
                    self.functions[id].unwrap()
                } else {
                    let id = namespace.get_data_id(name);
                    self.data_objects[id].unwrap()
                }
            }
            ir::ExternalName::LibCall(ref libcall) => {
                let name = (self.libcall_names)(*libcall);
                if let Some(symbol) = self.object.symbol_id(name.as_bytes()) {
                    symbol
                } else if let Some(symbol) = self.libcalls.get(libcall) {
                    *symbol
                } else {
                    let symbol = self.object.add_symbol(Symbol {
                        name: name.as_bytes().to_vec(),
                        value: 0,
                        size: 0,
                        kind: SymbolKind::Text,
                        scope: SymbolScope::Unknown,
                        weak: false,
                        section: SymbolSection::Undefined,
                        flags: SymbolFlags::None,
                    });
                    self.libcalls.insert(*libcall, symbol);
                    symbol
                }
            }
            _ => panic!("invalid ExternalName {}", name),
        }
    }
}

fn translate_linkage(linkage: Linkage) -> (SymbolScope, bool) {
    let scope = match linkage {
        Linkage::Import => SymbolScope::Unknown,
        Linkage::Local => SymbolScope::Compilation,
        Linkage::Export | Linkage::Preemptible => SymbolScope::Dynamic,
    };
    // TODO: this matches rustc_codegen_cranelift, but may be wrong.
    let weak = linkage == Linkage::Preemptible;
    (scope, weak)
}

#[derive(Clone)]
pub struct ObjectCompiledFunction {
    offset: u64,
    size: u32,
    section: SectionId,
    relocs: Vec<RelocRecord>,
}

#[derive(Clone)]
pub struct ObjectCompiledData {
    offset: u64,
    section: SectionId,
    relocs: Vec<RelocRecord>,
}

/// This is the output of `Module`'s
/// [`finish`](../cranelift_module/struct.Module.html#method.finish) function.
/// It contains the generated `Object` and other information produced during
/// compilation.
pub struct ObjectProduct {
    /// Object artifact with all functions and data from the module defined.
    pub object: Object,
    /// Symbol IDs for functions (both declared and defined).
    pub functions: SecondaryMap<FuncId, Option<SymbolId>>,
    /// Symbol IDs for data objects (both declared and defined).
    pub data_objects: SecondaryMap<DataId, Option<SymbolId>>,
    /// Trap sites for defined functions.
    pub traps: SecondaryMap<FuncId, Vec<ObjectTrapSite>>,
}

impl ObjectProduct {
    /// Return the `SymbolId` for the given function.
    #[inline]
    pub fn function_symbol(&self, id: FuncId) -> SymbolId {
        self.functions[id].unwrap()
    }

    /// Return the `SymbolId` for the given data object.
    #[inline]
    pub fn data_symbol(&self, id: DataId) -> SymbolId {
        self.data_objects[id].unwrap()
    }

    /// Write the object bytes in memory.
    #[inline]
    pub fn emit(self) -> Result<Vec<u8>, String> {
        self.object.write()
    }
}

#[derive(Clone)]
struct RelocRecord {
    offset: CodeOffset,
    name: ir::ExternalName,
    kind: RelocationKind,
    encoding: RelocationEncoding,
    size: u8,
    addend: Addend,
}

#[derive(Default)]
struct ObjectRelocSink {
    relocs: Vec<RelocRecord>,
}

impl RelocSink for ObjectRelocSink {
    fn reloc_ebb(&mut self, _offset: CodeOffset, _reloc: Reloc, _ebb_offset: CodeOffset) {
        unimplemented!();
    }

    fn reloc_external(
        &mut self,
        offset: CodeOffset,
        reloc: Reloc,
        name: &ir::ExternalName,
        addend: Addend,
    ) {
        let (kind, encoding, size) = match reloc {
            Reloc::Abs4 => (RelocationKind::Absolute, RelocationEncoding::Generic, 32),
            Reloc::Abs8 => (RelocationKind::Absolute, RelocationEncoding::Generic, 64),
            Reloc::X86PCRel4 => (RelocationKind::Relative, RelocationEncoding::Generic, 32),
            Reloc::X86CallPCRel4 => (RelocationKind::Relative, RelocationEncoding::X86Branch, 32),
            // TODO: Get Cranelift to tell us when we can use
            // R_X86_64_GOTPCRELX/R_X86_64_REX_GOTPCRELX.
            Reloc::X86CallPLTRel4 => (
                RelocationKind::PltRelative,
                RelocationEncoding::X86Branch,
                32,
            ),
            Reloc::X86GOTPCRel4 => (RelocationKind::GotRelative, RelocationEncoding::Generic, 32),
            // FIXME
            _ => unimplemented!(),
        };
        self.relocs.push(RelocRecord {
            offset,
            name: name.clone(),
            kind,
            encoding,
            size,
            addend,
        });
    }

    fn reloc_jt(&mut self, _offset: CodeOffset, reloc: Reloc, _jt: ir::JumpTable) {
        match reloc {
            Reloc::X86PCRelRodata4 => {
                // Not necessary to record this unless we are going to split apart code and its
                // jumptbl/rodata.
            }
            _ => {
                panic!("Unhandled reloc");
            }
        }
    }

    fn reloc_constant(&mut self, _offset: CodeOffset, reloc: Reloc, _jt: ir::ConstantOffset) {
        match reloc {
            Reloc::X86PCRelRodata4 => {
                // Not necessary to record this unless we are going to split apart code and its
                // jumptbl/rodata.
            }
            _ => {
                panic!("Unhandled reloc");
            }
        }
    }
}