From 003fd729899036096bb755c9e09124be1c09bd2a Mon Sep 17 00:00:00 2001
From: Wouter Nuijten <wouternuijten@gmail.com>
Date: Thu, 18 Apr 2024 11:44:31 +0200
Subject: [PATCH 1/9] Bug fix for duplicate arguments

---
 src/graph_engine.jl              |  8 +++++++-
 src/model_generator.jl           |  1 +
 test/graph_construction_tests.jl | 35 ++++++++++++++++++++++++++++----
 3 files changed, 39 insertions(+), 5 deletions(-)

diff --git a/src/graph_engine.jl b/src/graph_engine.jl
index 73a38139..e881f32b 100644
--- a/src/graph_engine.jl
+++ b/src/graph_engine.jl
@@ -1500,7 +1500,13 @@ function add_edge!(
     label = EdgeLabel(interface_name, index)
     neighbor_node_label = unroll(variable_node_id)
     addneighbor!(factor_node_propeties, neighbor_node_label, label, model[neighbor_node_label])
-    model.graph[unroll(variable_node_id), factor_node_id] = label
+    if MetaGraphsNext.add_edge!(model.graph, unroll(variable_node_id), factor_node_id, label)
+        return nothing
+    else
+        error(
+            lazy"Trying to create duplicate edge ($(unroll(variable_node_id)), $(factor_node_id)) while creating edge $(label) of factor node with functional form $(factor_node_propeties.fform)"
+        )
+    end
 end
 
 function add_edge!(
diff --git a/src/model_generator.jl b/src/model_generator.jl
index c0ad0d0a..0ac1b9e7 100644
--- a/src/model_generator.jl
+++ b/src/model_generator.jl
@@ -1,3 +1,4 @@
+using MetaGraphsNext.Graphs
 
 """
     ModelGenerator(model, kwargs, plugins)
diff --git a/test/graph_construction_tests.jl b/test/graph_construction_tests.jl
index 350960ed..e25835fb 100644
--- a/test/graph_construction_tests.jl
+++ b/test/graph_construction_tests.jl
@@ -1500,21 +1500,48 @@ end
     end
 end
 
-@testitem "LazyIndex should support empty indices if array is passed" begin 
+@testitem "LazyIndex should support empty indices if array is passed" begin
     import GraphPPL: create_model, getorcreate!, NodeCreationOptions, LazyIndex
 
     include("testutils.jl")
 
-    @model function foo(y) 
+    @model function foo(y)
         x ~ MvNormal([1, 1], [1 0.0; 0.0 1.0])
         y ~ MvNormal(x, [1.0 0.0; 0.0 1.0])
     end
 
-    model = create_model(foo()) do model, ctx 
-        return (; y = getorcreate!(model, ctx, NodeCreationOptions(kind = :data, factorized = true), :y, LazyIndex([ 1.0, 1.0 ])))
+    model = create_model(foo()) do model, ctx
+        return (; y = getorcreate!(model, ctx, NodeCreationOptions(kind = :data, factorized = true), :y, LazyIndex([1.0, 1.0])))
     end
 
     @test length(collect(filter(as_node(MvNormal), model))) == 2
     @test length(collect(filter(as_variable(:x), model))) == 1
     @test length(collect(filter(as_variable(:y), model))) == 1
+end
+
+@testitem "Node arguments must be unique" begin
+    import GraphPPL: create_model, getorcreate!, NodeCreationOptions, LazyIndex
+
+    include("testutils.jl")
+
+    @model function my_model(obs, N, sigma)
+        local p
+        for i in 1:N
+            p[i] ~ Beta(1, 1)
+        end
+        local x
+        for i in 1:N
+            x[i] ~ Bernoulli(p[i])
+        end
+        local C
+        for i in 1:N
+            C ~ C + x[i]
+        end
+        obs ~ NormalMeanVariance(C, sigma^2)
+    end
+
+    @test_throws ErrorException create_model(my_model(N = 3, sigma = 1.0)) do model, ctx
+        obs = getorcreate!(model, ctx, NodeCreationOptions(kind = :data, factorized = true), :obs, LazyIndex([0.0]))
+        return (obs = obs,)
+    end
 end
\ No newline at end of file

From b8a162113eef3a2fec8ccec1e4b46881e2b888e5 Mon Sep 17 00:00:00 2001
From: Wouter Nuijten <wouternuijten@gmail.com>
Date: Thu, 18 Apr 2024 12:56:57 +0200
Subject: [PATCH 2/9] write custom vertex adding function

---
 src/graph_engine.jl | 12 ++++++++++--
 1 file changed, 10 insertions(+), 2 deletions(-)

diff --git a/src/graph_engine.jl b/src/graph_engine.jl
index e881f32b..a5a46af2 100644
--- a/src/graph_engine.jl
+++ b/src/graph_engine.jl
@@ -949,6 +949,13 @@ Returns a collection of aliases for `fform` depending on the `backend`.
 aliases(backend, fform) = error("Backend $backend must implement a method for `aliases` for `$(fform)`.")
 aliases(model::Model, fform::F) where {F} = aliases(getbackend(model), fform)
 
+function add_vertex!(model::Model, label, data)
+    code = nv(model) + 1
+    model.graph.vertex_labels[code] = label
+    model.graph.vertex_properties[label] = (code, data)
+    Graphs.add_vertex!(model.graph.graph)
+end
+
 """
     copy_markov_blanket_to_child_context(child_context::Context, interfaces::NamedTuple)
 
@@ -1309,7 +1316,8 @@ function add_variable_node!(model::Model, context::Context, options::NodeCreatio
     )
 
     context[name, index] = label
-    model[label] = nodedata
+    add_vertex!(model, label, nodedata)
+    # model[label] = nodedata
 
     return label
 end
@@ -1428,7 +1436,7 @@ function add_atomic_factor_node!(model::Model, context::Context, options::NodeCr
         UnionPluginType(FactorNodePlugin(), FactorAndVariableNodesPlugin()), model, context, potential_label, potential_nodedata, options
     )
 
-    model[label] = nodedata
+    add_vertex!(model, label, nodedata)
     context[factornode_id] = label
 
     return label, nodedata, convert(FactorNodeProperties, getproperties(nodedata))

From 55109abe51f0528cdd24f0acfaf4fbb33059aad8 Mon Sep 17 00:00:00 2001
From: Wouter Nuijten <wouternuijten@gmail.com>
Date: Thu, 18 Apr 2024 13:37:44 +0200
Subject: [PATCH 3/9] Custom implementation for add vertex

---
 src/graph_engine.jl        | 12 +++++++++---
 test/graph_engine_tests.jl |  4 ++--
 2 files changed, 11 insertions(+), 5 deletions(-)

diff --git a/src/graph_engine.jl b/src/graph_engine.jl
index a5a46af2..a3a1cc75 100644
--- a/src/graph_engine.jl
+++ b/src/graph_engine.jl
@@ -956,6 +956,12 @@ function add_vertex!(model::Model, label, data)
     Graphs.add_vertex!(model.graph.graph)
 end
 
+function add_edge!(model::Model, src, dst, data)
+    code_src, code_dst = MetaGraphsNext.code_for(model.graph, src), MetaGraphsNext.code_for(model.graph, dst)
+    model.graph.edge_data[(src, dst)] = data
+    return Graphs.add_edge!(model.graph.graph, code_src, code_dst)
+end
+
 """
     copy_markov_blanket_to_child_context(child_context::Context, interfaces::NamedTuple)
 
@@ -1508,13 +1514,13 @@ function add_edge!(
     label = EdgeLabel(interface_name, index)
     neighbor_node_label = unroll(variable_node_id)
     addneighbor!(factor_node_propeties, neighbor_node_label, label, model[neighbor_node_label])
-    if MetaGraphsNext.add_edge!(model.graph, unroll(variable_node_id), factor_node_id, label)
-        return nothing
-    else
+    edge_added = add_edge!(model, unroll(variable_node_id), factor_node_id, label)
+    if !edge_added
         error(
             lazy"Trying to create duplicate edge ($(unroll(variable_node_id)), $(factor_node_id)) while creating edge $(label) of factor node with functional form $(factor_node_propeties.fform)"
         )
     end
+    return label
 end
 
 function add_edge!(
diff --git a/test/graph_engine_tests.jl b/test/graph_engine_tests.jl
index 45ed550f..5e88caa5 100644
--- a/test/graph_engine_tests.jl
+++ b/test/graph_engine_tests.jl
@@ -1842,13 +1842,13 @@ end
     model = create_test_model()
     ctx = getcontext(model)
     options = NodeCreationOptions()
-    x, xdata, xproperties = GraphPPL.add_atomic_factor_node!(model, ctx, options, sum)
     y = getorcreate!(model, ctx, :y, nothing)
 
     variable_nodes = [getorcreate!(model, ctx, i, nothing) for i in [:a, :b, :c]]
+    x, xdata, xproperties = GraphPPL.add_atomic_factor_node!(model, ctx, options, sum)
     add_edge!(model, x, xproperties, variable_nodes, :interface)
 
-    @test ne(model) == 3 && model[x, variable_nodes[1]] == EdgeLabel(:interface, 1)
+    @test ne(model) == 3 && model[variable_nodes[1], x] == EdgeLabel(:interface, 1)
 end
 
 @testitem "default_parametrization" begin

From 6d2fe8bff84e344fb73fa68e519eae13c443cf8e Mon Sep 17 00:00:00 2001
From: Wouter Nuijten <wouternuijten@gmail.com>
Date: Thu, 18 Apr 2024 13:48:34 +0200
Subject: [PATCH 4/9] Add additional tests

---
 test/graph_construction_tests.jl | 26 ++++++++++++++++++++++++--
 1 file changed, 24 insertions(+), 2 deletions(-)

diff --git a/test/graph_construction_tests.jl b/test/graph_construction_tests.jl
index e25835fb..60688039 100644
--- a/test/graph_construction_tests.jl
+++ b/test/graph_construction_tests.jl
@@ -1540,8 +1540,30 @@ end
         obs ~ NormalMeanVariance(C, sigma^2)
     end
 
-    @test_throws ErrorException create_model(my_model(N = 3, sigma = 1.0)) do model, ctx
-        obs = getorcreate!(model, ctx, NodeCreationOptions(kind = :data, factorized = true), :obs, LazyIndex([0.0]))
+    @test_throws "Trying to create duplicate edge" create_model(my_model(N = 3, sigma = 1.0)) do model, ctx
+        obs = getorcreate!(model, ctx, NodeCreationOptions(kind = :data, factorized = true), :obs, LazyIndex(0.0))
+        return (obs = obs,)
+    end
+
+    @model function my_model(obs, N, sigma)
+        local p
+        for i in 1:N
+            p[i] ~ Beta(1, 1)
+        end
+        local x
+        for i in 1:N
+            x[i] ~ Bernoulli(p[i])
+        end
+        accum_C = x[1]
+        for i in 2:N
+            next_C ~ accum_C + x[i]
+            accum_C = next_C
+        end
+        obs ~ NormalMeanVariance(accum_C, sigma^2)
+    end
+
+    @test_throws "Trying to create duplicate edge" create_model(my_model(N = 3, sigma = 1.0)) do model, ctx
+        obs = getorcreate!(model, ctx, NodeCreationOptions(kind = :data, factorized = true), :obs, LazyIndex(0.0))
         return (obs = obs,)
     end
 end
\ No newline at end of file

From 21959c93a98949e53ec310e2c71a0ed5efd0f631 Mon Sep 17 00:00:00 2001
From: Wouter Nuijten <wouternuijten@gmail.com>
Date: Thu, 18 Apr 2024 13:49:45 +0200
Subject: [PATCH 5/9] remove comment

---
 src/graph_engine.jl | 3 ---
 1 file changed, 3 deletions(-)

diff --git a/src/graph_engine.jl b/src/graph_engine.jl
index a3a1cc75..c091779c 100644
--- a/src/graph_engine.jl
+++ b/src/graph_engine.jl
@@ -1320,11 +1320,8 @@ function add_variable_node!(model::Model, context::Context, options::NodeCreatio
     label, nodedata = preprocess_plugins(
         UnionPluginType(VariableNodePlugin(), FactorAndVariableNodesPlugin()), model, context, potential_label, potential_nodedata, options
     )
-
     context[name, index] = label
     add_vertex!(model, label, nodedata)
-    # model[label] = nodedata
-
     return label
 end
 

From ad5844d93f556dca47ceacd6befd8fda45c7bf1a Mon Sep 17 00:00:00 2001
From: Wouter Nuijten <wouternuijten@gmail.com>
Date: Thu, 18 Apr 2024 13:50:27 +0200
Subject: [PATCH 6/9] Remove obsolete import

---
 src/model_generator.jl | 2 --
 1 file changed, 2 deletions(-)

diff --git a/src/model_generator.jl b/src/model_generator.jl
index 0ac1b9e7..8be780b5 100644
--- a/src/model_generator.jl
+++ b/src/model_generator.jl
@@ -1,5 +1,3 @@
-using MetaGraphsNext.Graphs
-
 """
     ModelGenerator(model, kwargs, plugins)
 

From 44af1c00c2cd913e296ad7c073b98eb02b321224 Mon Sep 17 00:00:00 2001
From: Wouter Nuijten <wouternuijten@gmail.com>
Date: Thu, 18 Apr 2024 13:50:50 +0200
Subject: [PATCH 7/9] Update model_generator.jl

---
 src/model_generator.jl | 1 +
 1 file changed, 1 insertion(+)

diff --git a/src/model_generator.jl b/src/model_generator.jl
index 8be780b5..c0ad0d0a 100644
--- a/src/model_generator.jl
+++ b/src/model_generator.jl
@@ -1,3 +1,4 @@
+
 """
     ModelGenerator(model, kwargs, plugins)
 

From 0701b46e1e825c5bcb99936d3a53835c5699e6c9 Mon Sep 17 00:00:00 2001
From: Bagaev Dmitry <bvdmitri@gmail.com>
Date: Thu, 18 Apr 2024 14:26:38 +0200
Subject: [PATCH 8/9] more tests to the god of tests

---
 src/graph_engine.jl              | 18 +++++++--
 test/graph_construction_tests.jl | 69 ++++++++++++++++++++++++++------
 test/graph_engine_tests.jl       | 10 +++++
 3 files changed, 80 insertions(+), 17 deletions(-)

diff --git a/src/graph_engine.jl b/src/graph_engine.jl
index c091779c..689a8a8a 100644
--- a/src/graph_engine.jl
+++ b/src/graph_engine.jl
@@ -962,6 +962,11 @@ function add_edge!(model::Model, src, dst, data)
     return Graphs.add_edge!(model.graph.graph, code_src, code_dst)
 end
 
+function has_edge(model::Model, src, dst)
+    code_src, code_dst = MetaGraphsNext.code_for(model.graph, src), MetaGraphsNext.code_for(model.graph, dst)
+    return Graphs.has_edge(model.graph.graph, code_src, code_dst)
+end
+
 """
     copy_markov_blanket_to_child_context(child_context::Context, interfaces::NamedTuple)
 
@@ -1511,11 +1516,16 @@ function add_edge!(
     label = EdgeLabel(interface_name, index)
     neighbor_node_label = unroll(variable_node_id)
     addneighbor!(factor_node_propeties, neighbor_node_label, label, model[neighbor_node_label])
-    edge_added = add_edge!(model, unroll(variable_node_id), factor_node_id, label)
+    edge_added = add_edge!(model, neighbor_node_label, factor_node_id, label)
     if !edge_added
-        error(
-            lazy"Trying to create duplicate edge ($(unroll(variable_node_id)), $(factor_node_id)) while creating edge $(label) of factor node with functional form $(factor_node_propeties.fform)"
-        )
+        # Double check if the edge has already been added
+        if has_edge(model, neighbor_node_label, factor_node_id)
+            error(
+                lazy"Trying to create duplicate edge $(label) between variable $(neighbor_node_label) and factor node $(factor_node_id). Make sure that all the arguments to the `~` operator are unique (both left hand side and right hand side)."
+            )
+        else
+            error(lazy"Cannot create an edge $(label) between variable $(neighbor_node_label) and factor node $(factor_node_id).")
+        end
     end
     return label
 end
diff --git a/test/graph_construction_tests.jl b/test/graph_construction_tests.jl
index 60688039..6ad26a29 100644
--- a/test/graph_construction_tests.jl
+++ b/test/graph_construction_tests.jl
@@ -1524,46 +1524,89 @@ end
 
     include("testutils.jl")
 
+    @model function simple_model_duplicate_1()
+        x ~ Normal(0.0, 1.0)
+        y ~ x + x
+    end
+
+    @model function simple_model_duplicate_2()
+        x ~ Normal(0.0, 1.0)
+        y ~ x + x + x
+    end
+
+    @model function simple_model_duplicate_3()
+        x ~ Normal(0.0, 1.0)
+        y ~ Normal(x, x)
+    end
+
+    @model function simple_model_duplicate_4()
+        x ~ Normal(0.0, 1.0)
+        hide_x = x
+        y ~ Normal(hide_x, x)
+    end
+
+    @model function simple_model_duplicate_5()
+        x ~ Normal(0.0, 1.0)
+        x ~ Normal(x, 1)
+    end
+
+    @model function simple_model_duplicate_6()
+        x ~ Normal(0.0, 1.0)
+        hide_x = x
+        hide_x ~ Normal(x, 1)
+    end
+
+    for modelfn in [
+        simple_model_duplicate_1,
+        simple_model_duplicate_2,
+        simple_model_duplicate_3,
+        simple_model_duplicate_4,
+        simple_model_duplicate_5,
+        simple_model_duplicate_6
+    ]
+        @test_throws r"Trying to create duplicate edge.*Make sure that all the arguments to the `~` operator are unique.*" create_model(
+            modelfn()
+        )
+    end
+
     @model function my_model(obs, N, sigma)
-        local p
-        for i in 1:N
-            p[i] ~ Beta(1, 1)
-        end
         local x
         for i in 1:N
-            x[i] ~ Bernoulli(p[i])
+            x[i] ~ Bernoulli(0.5)
         end
         local C
+        # This model creation is not allowed since `C` is used twice in the `~` operator
         for i in 1:N
             C ~ C + x[i]
         end
         obs ~ NormalMeanVariance(C, sigma^2)
     end
 
-    @test_throws "Trying to create duplicate edge" create_model(my_model(N = 3, sigma = 1.0)) do model, ctx
+    @test_throws r"Trying to create duplicate edge.*Make sure that all the arguments to the `~` operator are unique.*" create_model(
+        my_model(N = 3, sigma = 1.0)
+    ) do model, ctx
         obs = getorcreate!(model, ctx, NodeCreationOptions(kind = :data, factorized = true), :obs, LazyIndex(0.0))
         return (obs = obs,)
     end
 
     @model function my_model(obs, N, sigma)
-        local p
-        for i in 1:N
-            p[i] ~ Beta(1, 1)
-        end
         local x
         for i in 1:N
-            x[i] ~ Bernoulli(p[i])
+            x[i] ~ Bernoulli(0.5)
         end
         accum_C = x[1]
         for i in 2:N
+            # Here `next_C` will be used twice on the second iteration 
             next_C ~ accum_C + x[i]
             accum_C = next_C
         end
         obs ~ NormalMeanVariance(accum_C, sigma^2)
     end
 
-    @test_throws "Trying to create duplicate edge" create_model(my_model(N = 3, sigma = 1.0)) do model, ctx
+    @test_throws r"Trying to create duplicate edge.*Make sure that all the arguments to the `~` operator are unique.*" create_model(
+        my_model(N = 3, sigma = 1.0)
+    ) do model, ctx
         obs = getorcreate!(model, ctx, NodeCreationOptions(kind = :data, factorized = true), :obs, LazyIndex(0.0))
         return (obs = obs,)
     end
-end
\ No newline at end of file
+end
diff --git a/test/graph_engine_tests.jl b/test/graph_engine_tests.jl
index 5e88caa5..d4409e54 100644
--- a/test/graph_engine_tests.jl
+++ b/test/graph_engine_tests.jl
@@ -759,6 +759,7 @@ end
         EdgeLabel,
         getname,
         add_edge!,
+        has_edge,
         getproperties
 
     include("testutils.jl")
@@ -770,12 +771,21 @@ end
     b = NodeLabel(:b, 2)
     model[a] = NodeData(ctx, VariableNodeProperties(name = :a, index = nothing))
     model[b] = NodeData(ctx, FactorNodeProperties(fform = sum))
+    @test !has_edge(model, a, b)
+    @test !has_edge(model, b, a)
     add_edge!(model, b, getproperties(model[b]), a, :edge, 1)
+    @test has_edge(model, a, b)
+    @test has_edge(model, b, a)
     @test length(edges(model)) == 1
 
     c = NodeLabel(:c, 2)
     model[c] = NodeData(ctx, FactorNodeProperties(fform = sum))
+    @test !has_edge(model, a, c)
+    @test !has_edge(model, c, a)
     add_edge!(model, c, getproperties(model[c]), a, :edge, 2)
+    @test has_edge(model, a, c)
+    @test has_edge(model, c, a)
+
     @test length(edges(model)) == 2
 
     # Test 2: Test getting all edges from a model with a specific node

From 5b1ae25cc0088bb432ad98619ffe3ef5fe236433 Mon Sep 17 00:00:00 2001
From: Wouter Nuijten <wouternuijten@gmail.com>
Date: Thu, 18 Apr 2024 14:46:06 +0200
Subject: [PATCH 9/9] Add comments to unsafe functions

---
 src/graph_engine.jl | 2 ++
 1 file changed, 2 insertions(+)

diff --git a/src/graph_engine.jl b/src/graph_engine.jl
index 689a8a8a..1d19a9cd 100644
--- a/src/graph_engine.jl
+++ b/src/graph_engine.jl
@@ -950,6 +950,7 @@ aliases(backend, fform) = error("Backend $backend must implement a method for `a
 aliases(model::Model, fform::F) where {F} = aliases(getbackend(model), fform)
 
 function add_vertex!(model::Model, label, data)
+    # This is an unsafe procedure that implements behaviour from `MetaGraphsNext`. 
     code = nv(model) + 1
     model.graph.vertex_labels[code] = label
     model.graph.vertex_properties[label] = (code, data)
@@ -957,6 +958,7 @@ function add_vertex!(model::Model, label, data)
 end
 
 function add_edge!(model::Model, src, dst, data)
+    # This is an unsafe procedure that implements behaviour from `MetaGraphsNext`. 
     code_src, code_dst = MetaGraphsNext.code_for(model.graph, src), MetaGraphsNext.code_for(model.graph, dst)
     model.graph.edge_data[(src, dst)] = data
     return Graphs.add_edge!(model.graph.graph, code_src, code_dst)