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