feat: Generate a custom GOT array and save/load its layout

lewis-revill · CohenArthur · commit d9c456136119 · 2024-06-12T15:54:15.000+02:00
This commit contains changes that are required to generate a global
variable which will be used to store an array of addresses of other
global variables in the program. The location of the globals within this
array should be stable after a recompilation to support online change,
so we support loading a pre-existing layout to ensure that for all
globals that we see in both the current program and the pre-existing
layout, their positions remain the same. Otherwise, any new globals will
be placed in any empty spaces left by old globals, or appended on to the
end of the array. The layout will then be saved back the the file used
for saving and loading.

Currently, to use this feature the flag `--got-layout-file=&lt;file&gt;` must
be provided, which should specify the name of either a TOML or JSON file
to use to save, and optionally load if the file already exists, the GOT
layout. In future we will integrate this with a generic online change
flag, whereby it will not be necessary to ask for a GOT layout file when
we already know that we need it for online change.
diff --git a/compiler/plc_driver/src/cli.rs b/compiler/plc_driver/src/cli.rs
@@ -100,6 +100,18 @@ pub struct CompileParameters {
     ) ]
     pub hardware_config: Option<String>,
 
+    #[clap(
+        name = "got-layout-file",
+        long,
+        global = true,
+        help = "Obtain information about the current custom GOT layout from the given file if it exists.
+    Save information about the generated custom GOT layout to the given file.
+    Format is detected by extension.
+    Supported formats : json, toml",
+    parse(try_from_str = validate_config)
+    ) ]
+    pub got_layout_file: Option<String>,
+
     #[clap(
         name = "optimization",
         long,
@@ -379,6 +391,10 @@ impl CompileParameters {
         self.hardware_config.as_deref().and_then(get_config_format)
     }
 
+    pub fn got_layout_format(&self) -> Option<ConfigFormat> {
+        self.got_layout_file.as_deref().and_then(get_config_format)
+    }
+
     /// Returns the location where the build artifacts should be stored / output
     pub fn get_build_location(&self) -> Option<PathBuf> {
         match &self.commands {
diff --git a/compiler/plc_driver/src/lib.rs b/compiler/plc_driver/src/lib.rs
@@ -19,8 +19,8 @@ use std::{
 use cli::{CompileParameters, ParameterError, SubCommands};
 use pipelines::AnnotatedProject;
 use plc::{
-    codegen::CodegenContext, linker::LinkerType, output::FormatOption, DebugLevel, ErrorFormat,
-    OptimizationLevel, Target, Threads,
+    codegen::CodegenContext, linker::LinkerType, output::FormatOption, ConfigFormat, DebugLevel,
+    ErrorFormat, OptimizationLevel, Target, Threads,
 };
 
 use plc_diagnostics::{diagnostician::Diagnostician, diagnostics::Diagnostic};
@@ -50,6 +50,8 @@ pub struct CompileOptions {
     /// The name of the resulting compiled file
     pub output: String,
     pub output_format: FormatOption,
+    pub got_layout_file: Option<String>,
+    pub got_layout_format: Option<ConfigFormat>,
     pub optimization: OptimizationLevel,
     pub error_format: ErrorFormat,
     pub debug_level: DebugLevel,
@@ -63,6 +65,8 @@ impl Default for CompileOptions {
             build_location: None,
             output: String::new(),
             output_format: Default::default(),
+            got_layout_file: None,
+            got_layout_format: None,
             optimization: OptimizationLevel::None,
             error_format: ErrorFormat::None,
             debug_level: DebugLevel::None,
@@ -172,6 +176,8 @@ pub fn get_compilation_context<T: AsRef<str> + AsRef<OsStr> + Debug>(
         build_location: compile_parameters.get_build_location(),
         output: project.get_output_name(),
         output_format,
+        got_layout_file: compile_parameters.got_layout_file.clone(),
+        got_layout_format: compile_parameters.got_layout_format(),
         optimization: compile_parameters.optimization,
         error_format: compile_parameters.error_format,
         debug_level: compile_parameters.debug_level(),
diff --git a/compiler/plc_driver/src/pipelines.rs b/compiler/plc_driver/src/pipelines.rs
@@ -274,6 +274,7 @@ impl<T: SourceContainer + Sync> AnnotatedProject<T> {
             context,
             compile_options.root.as_deref(),
             &unit.file_name,
+            compile_options.got_layout_file.clone().zip(compile_options.got_layout_format),
             compile_options.optimization,
             compile_options.debug_level,
         );
diff --git a/src/codegen.rs b/src/codegen.rs
@@ -19,7 +19,7 @@ use self::{
 use crate::{
     output::FormatOption,
     resolver::{AstAnnotations, Dependency, StringLiterals},
-    DebugLevel, OptimizationLevel, Target,
+    ConfigFormat, DebugLevel, OptimizationLevel, Target,
 };
 
 use super::index::*;
@@ -71,6 +71,8 @@ pub struct CodeGen<'ink> {
     /// the debugging module creates debug information at appropriate locations
     pub debug: DebugBuilderEnum<'ink>,
 
+    pub got_layout_file: Option<(String, ConfigFormat)>,
+
     pub module_location: String,
 }
 
@@ -88,13 +90,14 @@ impl<'ink> CodeGen<'ink> {
         context: &'ink CodegenContext,
         root: Option<&Path>,
         module_location: &str,
+        got_layout_file: Option<(String, ConfigFormat)>,
         optimization_level: OptimizationLevel,
         debug_level: DebugLevel,
     ) -> CodeGen<'ink> {
         let module = context.create_module(module_location);
         module.set_source_file_name(module_location);
         let debug = debug::DebugBuilderEnum::new(context, &module, root, optimization_level, debug_level);
-        CodeGen { module, debug, module_location: module_location.to_string() }
+        CodeGen { module, debug, got_layout_file, module_location: module_location.to_string() }
     }
 
     pub fn generate_llvm_index(
@@ -117,8 +120,15 @@ impl<'ink> CodeGen<'ink> {
         )?;
         index.merge(llvm_type_index);
 
-        let mut variable_generator =
-            VariableGenerator::new(&self.module, &llvm, global_index, annotations, &index, &mut self.debug);
+        let mut variable_generator = VariableGenerator::new(
+            &self.module,
+            &llvm,
+            global_index,
+            annotations,
+            &index,
+            &mut self.debug,
+            self.got_layout_file.clone(),
+        );
 
         //Generate global variables
         let llvm_gv_index =
diff --git a/src/codegen/generators/variable_generator.rs b/src/codegen/generators/variable_generator.rs
@@ -5,11 +5,15 @@ use crate::{
     codegen::{debug::Debug, llvm_index::LlvmTypedIndex, llvm_typesystem::cast_if_needed},
     index::{get_initializer_name, Index, PouIndexEntry, VariableIndexEntry},
     resolver::{AnnotationMap, AstAnnotations, Dependency},
+    ConfigFormat,
 };
-
-use inkwell::{module::Module, values::GlobalValue};
+use indexmap::IndexSet;
+use inkwell::{module::Module, types::BasicTypeEnum, values::GlobalValue};
 use plc_ast::ast::LinkageType;
 use plc_diagnostics::diagnostics::Diagnostic;
+use std::collections::HashMap;
+use std::fs::{read_to_string, write};
+use std::path::Path;
 
 use super::{
     data_type_generator::get_default_for,
@@ -20,13 +24,48 @@ use super::{
 use crate::codegen::debug::DebugBuilderEnum;
 use crate::index::FxIndexSet;
 
+pub fn read_got_layout(location: &str, format: ConfigFormat) -> Result<HashMap<String, u64>, Diagnostic> {
+    if !Path::new(location).is_file() {
+        // Assume if the file doesn't exist that there is no existing GOT layout yet. write_got_layout will handle
+        // creating our file when we want to.
+        return Ok(HashMap::new());
+    }
+
+    let s =
+        read_to_string(location).map_err(|_| Diagnostic::new("GOT layout could not be read from file"))?;
+    match format {
+        ConfigFormat::JSON => serde_json::from_str(&s)
+            .map_err(|_| Diagnostic::new("Could not deserialize GOT layout from JSON")),
+        ConfigFormat::TOML => {
+            toml::de::from_str(&s).map_err(|_| Diagnostic::new("Could not deserialize GOT layout from TOML"))
+        }
+    }
+}
+
+pub fn write_got_layout(
+    got_entries: HashMap<String, u64>,
+    location: &str,
+    format: ConfigFormat,
+) -> Result<(), Diagnostic> {
+    let s = match format {
+        ConfigFormat::JSON => serde_json::to_string(&got_entries)
+            .map_err(|_| Diagnostic::new("Could not serialize GOT layout to JSON"))?,
+        ConfigFormat::TOML => toml::ser::to_string(&got_entries)
+            .map_err(|_| Diagnostic::new("Could not serialize GOT layout to TOML"))?,
+    };
+
+    write(location, s).map_err(|_| Diagnostic::new("GOT layout could not be written to file"))?;
+    Ok(())
+}
+
 pub struct VariableGenerator<'ctx, 'b> {
     module: &'b Module<'ctx>,
     llvm: &'b Llvm<'ctx>,
     global_index: &'b Index,
     annotations: &'b AstAnnotations,
     types_index: &'b LlvmTypedIndex<'ctx>,
     debug: &'b mut DebugBuilderEnum<'ctx>,
+    got_layout_file: Option<(String, ConfigFormat)>,
 }
 
 impl<'ctx, 'b> VariableGenerator<'ctx, 'b> {
@@ -37,8 +76,9 @@ impl<'ctx, 'b> VariableGenerator<'ctx, 'b> {
         annotations: &'b AstAnnotations,
         types_index: &'b LlvmTypedIndex<'ctx>,
         debug: &'b mut DebugBuilderEnum<'ctx>,
+        got_layout_file: Option<(String, ConfigFormat)>,
     ) -> Self {
-        VariableGenerator { module, llvm, global_index, annotations, types_index, debug }
+        VariableGenerator { module, llvm, global_index, annotations, types_index, debug, got_layout_file }
     }
 
     pub fn generate_global_variables(
@@ -76,7 +116,7 @@ impl<'ctx, 'b> VariableGenerator<'ctx, 'b> {
             }
         });
 
-        for (name, variable) in globals {
+        for (name, variable) in &globals {
             let linkage =
                 if !variable.is_in_unit(location) { LinkageType::External } else { variable.get_linkage() };
             let global_variable = self.generate_global_variable(variable, linkage).map_err(|err| {
@@ -100,6 +140,46 @@ impl<'ctx, 'b> VariableGenerator<'ctx, 'b> {
             );
         }
 
+        if let Some((location, format)) = &self.got_layout_file {
+            let got_entries = read_got_layout(location.as_str(), *format)?;
+            let mut new_globals = Vec::new();
+            let mut new_got_entries = HashMap::new();
+            let mut new_got = HashMap::new();
+
+            for (name, _) in &globals {
+                if let Some(idx) = got_entries.get(&name.to_string()) {
+                    new_got_entries.insert(name.to_string(), *idx);
+                    new_got.insert(*idx, name.to_string());
+                } else {
+                    new_globals.push(name.to_string());
+                }
+            }
+
+            // Put any globals that weren't there last time in any free space in the GOT.
+            let mut idx: u64 = 0;
+            for name in &new_globals {
+                while new_got.contains_key(&idx) {
+                    idx += 1;
+                }
+                new_got_entries.insert(name.to_string(), idx);
+                new_got.insert(idx, name.to_string());
+            }
+
+            // Now we can write new_got_entries back out to a file.
+            write_got_layout(new_got_entries, location.as_str(), *format)?;
+
+            // Construct our GOT as a new global array. We initialise this array in the loader code.
+            let got_size = new_got.keys().max().map_or(0, |m| *m + 1);
+            let _got = self.llvm.create_global_variable(
+                self.module,
+                "__custom_got",
+                BasicTypeEnum::ArrayType(Llvm::get_array_type(
+                    BasicTypeEnum::PointerType(self.llvm.context.i8_type().ptr_type(0.into())),
+                    got_size.try_into().expect("the computed custom GOT size is too large"),
+                )),
+            );
+        }
+
         Ok(index)
     }
 
