Try #316:

Author: bors[bot]

src/compiler.jl

@@ -470,10 +470,7 @@ Builds the output expression.
 """
 function build_output(modelinfo, linenumbernode)
     ## Build the anonymous evaluator from the user-provided model definition.
-
-    # Remove the name.
     evaluatordef = deepcopy(modelinfo[:modeldef])
-    delete!(evaluatordef, :name)
 
     # Add the internal arguments to the user-specified arguments (positional + keywords).
     evaluatordef[:args] = vcat(
@@ -489,7 +486,13 @@ function build_output(modelinfo, linenumbernode)
     evaluatordef[:kwargs] = []
 
     # Replace the user-provided function body with the version created by DynamicPPL.
-    evaluatordef[:body] = modelinfo[:body]
+    # We use `MacroTools.@q begin ... end` instead of regular `quote ... end` to ensure
+    # that no new `LineNumberNode`s are added apart from the reference `linenumbernode`
+    # to the call site
+    evaluatordef[:body] = MacroTools.@q begin
+        $(linenumbernode)
+        $(modelinfo[:body])
+    end
 
     ## Build the model function.
 
@@ -498,24 +501,24 @@ function build_output(modelinfo, linenumbernode)
     defaults_namedtuple = modelinfo[:defaults_namedtuple]
 
     # Update the function body of the user-specified model.
-    # We use a name for the anonymous evaluator that does not conflict with other variables.
-    modeldef = modelinfo[:modeldef]
-    @gensym evaluator
     # We use `MacroTools.@q begin ... end` instead of regular `quote ... end` to ensure
     # that no new `LineNumberNode`s are added apart from the reference `linenumbernode`
     # to the call site
+    modeldef = modelinfo[:modeldef]
     modeldef[:body] = MacroTools.@q begin
         $(linenumbernode)
-        $evaluator = $(MacroTools.combinedef(evaluatordef))
         return $(DynamicPPL.Model)(
             $(QuoteNode(modeldef[:name])),
-            $evaluator,
+            $(modeldef[:name]),
             $allargs_namedtuple,
             $defaults_namedtuple,
         )
     end
 
-    return :($(Base).@__doc__ $(MacroTools.combinedef(modeldef)))
+    return MacroTools.@q begin
+        $(MacroTools.combinedef(evaluatordef))
+        $(Base).@__doc__ $(MacroTools.combinedef(modeldef))
+    end
 end
 
 function warn_empty(body)

src/model.jl

@@ -121,7 +121,7 @@ julia> @model function demo()
            x ~ Normal(m, 1)
            return (; m=m, x=x)
        end
-demo (generic function with 1 method)
+demo (generic function with 2 methods)
 
 julia> model = demo();
 
@@ -161,7 +161,7 @@ julia> @model function demo_mv(::Type{TV}=Float64) where {TV}
            m[2] ~ Normal()
            return m
        end
-demo_mv (generic function with 2 methods)
+demo_mv (generic function with 3 methods)
 
 julia> model = demo_mv();
 
@@ -192,13 +192,13 @@ the use of [`@submodel`](@ref).
 
 ```jldoctest condition
 julia> @model demo_inner() = m ~ Normal()
-demo_inner (generic function with 1 method)
+demo_inner (generic function with 2 methods)
 
 julia> @model function demo_outer()
            m = @submodel demo_inner()
            return m
        end
-demo_outer (generic function with 1 method)
+demo_outer (generic function with 2 methods)
 
 julia> model = demo_outer();
 
@@ -218,7 +218,7 @@ julia> @model function demo_outer_prefix()
            m = @submodel inner demo_inner()
            return m
        end
-demo_outer_prefix (generic function with 1 method)
+demo_outer_prefix (generic function with 2 methods)
 
 julia> # This doesn't work now!
        conditioned_model = demo_outer_prefix() | (m = 1.0, );
@@ -279,7 +279,7 @@ julia> @model function demo()
            x ~ Normal(m, 1)
            return (; m=m, x=x)
        end
-demo (generic function with 1 method)
+demo (generic function with 2 methods)
 
 julia> conditioned_model = condition(demo(), m = 1.0, x = 10.0);
 
@@ -333,7 +333,7 @@ julia> @model function demo()
            m ~ Normal()
            x ~ Normal(m, 1)
        end
-demo (generic function with 1 method)
+demo (generic function with 2 methods)
 
 julia> m = demo();
 
@@ -613,7 +613,7 @@ julia> @model function demo(xs)
            end
            return (m, )
        end
-demo (generic function with 1 method)
+demo (generic function with 2 methods)
 
 julia> model = demo(randn(10));
 

test/compiler.jl

@@ -39,13 +39,15 @@ end
 
             return x, y
         end
+        @test length(methods(testmodel_comp)) == 2
         testmodel_comp(1.0, 1.2)
 
         # check if drawing from the prior works
         @model function testmodel01(x=missing)
             x ~ Normal()
             return x
         end
+        @test length(methods(testmodel01)) == 3
         f0_mm = testmodel01()
         @test mean(f0_mm() for _ in 1:1000) ≈ 0.0 atol = 0.1
 
@@ -58,6 +60,7 @@ end
             x[2] ~ Normal()
             return x
         end
+        @test length(methods(testmodel02)) == 3
         f0_mm = testmodel02()
         @test all(x -> isapprox(x, 0; atol=0.1), mean(f0_mm() for _ in 1:1000))
 
@@ -66,6 +69,7 @@ end
             return x
         end
         f01_mm = testmodel03()
+        @test length(methods(testmodel03)) == 3
         @test mean(f01_mm() for _ in 1:1000) ≈ 0.5 atol = 0.1
 
         # test if we get the correct return values
@@ -78,6 +82,7 @@ end
 
             return x1, x2
         end
+        @test length(methods(testmodel1)) == 2
         f1_mm = testmodel1(1.0, 10.0)
         @test f1_mm() == (1, 10)
 
@@ -95,6 +100,7 @@ end
 
             return x1, x2
         end
+        @test length(methods(testmodel2)) == 2
         f1_mm = testmodel2(; x1=1.0, x2=10.0)
         @test f1_mm() == (1, 10)
 
@@ -461,4 +467,31 @@ end
             model = @model(x -> (x ~ Normal()))
         end
     end
+
+    @testset "dispatching with model" begin
+        f(x) = false
+
+        @model demo() = x ~ Normal()
+        @test !f(demo())
+        f(::Model{typeof(demo)}) = true
+        @test f(demo())
+
+        # Leads to re-definition of `demo` and trait is not affected.
+        @test length(methods(demo)) == 2
+        @model demo() = x ~ Normal()
+        @test length(methods(demo)) == 2
+        @test f(demo())
+
+        # Ensure we can specialize on arguments.
+        @model demo(x) = x ~ Normal()
+        length(methods(demo))
+        @test f(demo(1.0))
+        f(::Model{typeof(demo),(:x,)}) = false
+        @test !f(demo(1.0))
+        @test f(demo()) # should still be `true`
+
+        # Set it to `false` again.
+        f(::Model{typeof(demo),()}) = false
+        @test !f(demo())
+    end
 end