diff --git a/src/ir/display.rs b/src/ir/display.rs
index c5bfecd..fb53a40 100644
--- a/src/ir/display.rs
+++ b/src/ir/display.rs
@@ -322,14 +322,16 @@ impl<'a, PD: PrintDecorator> Display for ModuleDisplay<'a, PD> {
                         sig_strs.get(&sig).unwrap()
                     )?;
 
-                    if let Some(decorator) = &self.decorators {
-                        let decorator = &(*decorator)(func);
-                        writeln!(
-                            f,
-                            "{}",
-                            body.display_with_decorator("    ", Some(self.module), decorator)
-                        )?;
-                    }
+                    let decorator = self
+                        .decorators
+                        .as_ref()
+                        .map(|decorator_fn| decorator_fn(func));
+
+                    writeln!(
+                        f,
+                        "{}",
+                        body.display_with_decorator("    ", Some(self.module), decorator.as_ref())
+                    )?;
                 }
                 FuncDecl::Lazy(sig, name, reader) => {
                     writeln!(
diff --git a/src/ir/func.rs b/src/ir/func.rs
index f2f769c..0dccfba 100644
--- a/src/ir/func.rs
+++ b/src/ir/func.rs
@@ -471,14 +471,14 @@ impl FunctionBody {
         &'a self,
         indent: &'a str,
         module: Option<&'a Module>,
-        decorator: &'a PD,
+        decorator: Option<&'a PD>,
     ) -> FunctionBodyDisplay<'a, PD> {
         FunctionBodyDisplay {
             body: self,
             indent,
             verbose: false,
             module,
-            decorator: Some(&decorator),
+            decorator: decorator,
         }
     }