From 6b74e59835266d0a0013f7300f382fcd51d9dd87 Mon Sep 17 00:00:00 2001
From: Pedro Maciel Xavier <pmacielx@purdue.edu>
Date: Thu, 19 Sep 2024 09:32:58 -0400
Subject: [PATCH] Add attributes to access inner models

---
 src/attributes/compiler.jl    | 24 +++++++++++++++++++++++-
 test/integration/interface.jl |  3 +++
 2 files changed, 26 insertions(+), 1 deletion(-)

diff --git a/src/attributes/compiler.jl b/src/attributes/compiler.jl
index b1775c2..a3d04ce 100644
--- a/src/attributes/compiler.jl
+++ b/src/attributes/compiler.jl
@@ -8,7 +8,7 @@ const CI   = MOI.ConstraintIndex
 import PseudoBooleanOptimization as PBO
 import QUBOTools
 
-import ..ToQUBO: Optimizer
+import ..ToQUBO: Optimizer, PreQUBOModel, QUBOModel
 import ..Encoding
 import ..Virtual
 
@@ -20,6 +20,28 @@ abstract type CompilerAttribute <: MOI.AbstractOptimizerAttribute end
 
 MOI.supports(::Optimizer, ::A) where {A<:CompilerAttribute} = true
 
+@doc raw"""
+    SourceModel()
+"""
+struct SourceModel <: CompilerAttribute end
+
+MOI.is_set_by_optimize(::SourceModel) = true
+
+function MOI.get(model::Optimizer{T}, ::SourceModel)::PreQUBOModel{T} where {T}
+    return model.source_model
+end
+
+@doc raw"""
+    TargetModel()
+"""
+struct TargetModel <: CompilerAttribute end
+
+MOI.is_set_by_optimize(::TargetModel) = true
+
+function MOI.get(model::Optimizer{T}, ::TargetModel)::QUBOModel{T} where {T}
+    return model.target_model
+end
+
 @doc raw"""
     CompilationTime()
 """
diff --git a/test/integration/interface.jl b/test/integration/interface.jl
index aad9f3d..9f54376 100644
--- a/test/integration/interface.jl
+++ b/test/integration/interface.jl
@@ -170,6 +170,9 @@ function test_interface_moi()
                 @test MOI.get(model, MOI.RawOptimizerAttribute("num_reads")) == 1_001
 
                 # ToQUBO Attributes
+                @test MOI.get(model, Attributes.SourceModel()) isa ToQUBO.PreQUBOModel{Float64}
+                @test MOI.get(model, Attributes.TargetModel()) isa ToQUBO.QUBOModel{Float64}
+
                 @test MOI.get(model, Attributes.Architecture()) isa QUBOTools.GenericArchitecture
                 MOI.set(model, Attributes.Architecture(), SuperArchitecture(true))
                 @test MOI.get(model, Attributes.Architecture()) isa SuperArchitecture