From 9f392a7a900a685f60e57466d8422c7a2a91608a Mon Sep 17 00:00:00 2001
From: Tor Erlend Fjelde <tor.erlend95@gmail.com>
Date: Thu, 31 Oct 2024 12:01:43 +0100
Subject: [PATCH] added InteractiveUtils.jl as a dep + `model_typed` and
 `model_warntype` for the checking of the model's evaluator

---
 Project.toml       |  1 +
 src/debug_utils.jl | 33 +++++++++++++++++++++++++++++++++
 2 files changed, 34 insertions(+)

diff --git a/Project.toml b/Project.toml
index 92fb67ddd..879feeff1 100644
--- a/Project.toml
+++ b/Project.toml
@@ -14,6 +14,7 @@ Compat = "34da2185-b29b-5c13-b0c7-acf172513d20"
 ConstructionBase = "187b0558-2788-49d3-abe0-74a17ed4e7c9"
 Distributions = "31c24e10-a181-5473-b8eb-7969acd0382f"
 DocStringExtensions = "ffbed154-4ef7-542d-bbb7-c09d3a79fcae"
+InteractiveUtils = "b77e0a4c-d291-57a0-90e8-8db25a27a240"
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
 LogDensityProblems = "6fdf6af0-433a-55f7-b3ed-c6c6e0b8df7c"
 LogDensityProblemsAD = "996a588d-648d-4e1f-a8f0-a84b347e47b1"
diff --git a/src/debug_utils.jl b/src/debug_utils.jl
index a614508bf..873456a44 100644
--- a/src/debug_utils.jl
+++ b/src/debug_utils.jl
@@ -5,6 +5,7 @@ using ..DynamicPPL: broadcast_safe, AbstractContext, childcontext
 
 using Random: Random
 using Accessors: Accessors
+using InteractiveUtils: InteractiveUtils
 
 using DocStringExtensions
 using Distributions
@@ -678,4 +679,36 @@ function has_static_constraints(
     return all_the_same(transforms)
 end
 
+"""
+    model_warntype(model[, varinfo, context])
+
+Check the type stability of the model's evaluator, warning about any potential issues.
+
+This simply calls `@code_warntype` on the model's evaluator, filling in internal arguments where needed.
+"""
+function model_warntype(model::Model, varinfo::AbstractVarInfo=VarInfo(model), context::AbstractContext=DefaultContext())
+    args, kwargs = DynamicPPL.make_evaluate_args_and_kwargs(model, varinfo, context)
+    return if isempty(kwargs)
+        InteractiveUtils.@code_warntype model.f(args...)
+    else
+        InteractiveUtils.@code_warntype model.f(args...; kwargs...)
+    end
+end
+
+"""
+    model_typed(model[, varinfo, context])
+
+Return the type inference for the model's evaluator.
+
+This simply calls `@code_typed` on the model's evaluator, filling in internal arguments where needed.
+"""
+function model_typed(model::Model, varinfo::AbstractVarInfo=VarInfo(model), context::AbstractContext=DefaultContext())
+    args, kwargs = DynamicPPL.make_evaluate_args_and_kwargs(model, varinfo, context)
+    return if isempty(kwargs)
+        InteractiveUtils.@code_typed model.f(args...)
+    else
+        InteractiveUtils.@code_typed model.f(args...; kwargs...)
+    end
+end
+
 end