update: adding codescan analysis (#3)

.github/workflows/codescan-analysis.yml

@@ -0,0 +1,25 @@
+name: Semgrep Codescan
+
+on:
+  pull_request:
+    branches:
+      - main
+      - feature/**
+
+jobs:
+  semgrep-codescan:
+    name: Semgrep Codescan
+    runs-on: ubuntu-latest
+    steps:
+    - uses: actions/checkout@v4
+      with:
+        submodules: true
+    - uses: actions/setup-python@v5
+      with:
+        python-version: '3.10'
+    - name: Install Semgrep
+      run: python3 -m pip install semgrep
+    - name: Get rules from JuliaComputing
+      run: git clone https://github.com/JuliaComputing/semgrep-rules-julia.git
+    - name: Run Semgrep Julia rules
+      run: semgrep --error --config semgrep-rules-julia/rules .

ext/BraketSimulatorPythonExt/BraketSimulatorPythonExt.jl

@@ -254,7 +254,9 @@ function BraketSimulator.parse_program(simulator::D, program::OpenQasmProgram, s
     if shots > 0
         py_circ.instructions += py_circ.basis_rotation_instructions
     end
-    return ir(pyconvert(Circuit, py_circ), Val(:JAQCD))
+    circ = pyconvert(Circuit, py_circ)
+    @assert qubit_count(circ) <= properties(simulator).paradigm.qubitCount "parsed circuit's qubit count $(qubit_count(circ)) is larger than maximum ($(properties(simulator).paradigm.qubitCount)) for simulator type $D."
+    return ir(circ, Val(:JAQCD))
 end
 
 function simulate(

src/BraketSimulator.jl

@@ -145,13 +145,12 @@ function _generate_results(
     result_types::Vector,
     simulator::D,
 ) where {D<:AbstractSimulator}
-    result_values = [calculate(result_type, simulator) for result_type in result_types]
+    result_values = map(result_type -> calculate(result_type, simulator), result_types)
     result_values =
         [val isa Matrix ? Braket.complex_matrix_to_ir(val) : val for val in result_values]
-    return [
-        Braket.ResultTypeValue(result, result_value) for
-        (result, result_value) in zip(results, result_values)
-    ]
+    return map(zip(results, result_values)) do (result, result_value)
+        Braket.ResultTypeValue(result, result_value)
+    end
 end
 
 _translate_result_type(r::Braket.IR.Amplitude, qc::Int)     = Braket.Amplitude(r.states)
@@ -182,10 +181,10 @@ function _translate_result_types(
     results::Vector{Braket.AbstractProgramResult},
     qubit_count::Int,
 )
-    return [_translate_result_type(r, qubit_count) for r in results]
+    return map(result->_translate_result_type(result, qubit_count), results)
 end
 
-function _compute_exact_results(d::AbstractSimulator, program::Program, qc::Int, inputs::Dict{String, Float64})
+function _compute_exact_results(d::AbstractSimulator, program::Program, qc::Int)
     result_types = _translate_result_types(program.results, qc)
     _validate_result_types_qubits_exist(result_types, qc)
     return _generate_results(program.results, result_types, d)
@@ -225,7 +224,7 @@ function simulate(
         simulator = evolve!(simulator, operations)
     end
     @debug "Time for evolution: $(stats.time)"
-    results = shots == 0 && !isempty(program.results) ? _compute_exact_results(simulator, program, n_qubits, inputs) : [Braket.ResultTypeValue(result_type, 0.0) for result_type in program.results]
+    results = shots == 0 && !isempty(program.results) ? _compute_exact_results(simulator, program, n_qubits) : [Braket.ResultTypeValue(result_type, 0.0) for result_type in program.results]
     measured_qubits = get(kwargs, :measured_qubits, collect(0:n_qubits-1))
     isempty(measured_qubits) && (measured_qubits = collect(0:n_qubits-1))
     stats   = @timed _bundle_results(results, circuit_ir, simulator; measured_qubits=measured_qubits)
@@ -257,7 +256,7 @@ function simulate(
         simulator = evolve!(simulator, operations)
     end
     @debug "Time for evolution: $(stats.time)"
-    results = shots == 0 && !isempty(circuit_ir.results) ? _compute_exact_results(simulator, circuit_ir, qubit_count, inputs) : Braket.ResultTypeValue[]
+    results = shots == 0 && !isempty(circuit_ir.results) ? _compute_exact_results(simulator, circuit_ir, qubit_count) : Braket.ResultTypeValue[]
     measured_qubits = get(kwargs, :measured_qubits, collect(0:qubit_count-1))
     isempty(measured_qubits) && (measured_qubits = collect(0:qubit_count-1))
     stats = @timed _bundle_results(results, circuit_ir, simulator; measured_qubits=measured_qubits)

src/custom_gates.jl

@@ -62,6 +62,23 @@ Braket.chars(::Type{MultiQubitPhaseShift}) = "GPhase(ang)"
 Braket.qubit_count(g::MultiQubitPhaseShift{N}) where {N} = N
 Base.inv(g::MultiQubitPhaseShift{N}) where {N} = MultiQubitPhaseShift{N}((-g.angle[1],))
 
+function apply_gate!(
+    factor::ComplexF64,
+    diagonal::Braket.PauliEigenvalues{N},
+    state_vec::StateVector{T},
+    ts::Vararg{Int,N},
+) where {T<:Complex,N}
+    g_mat  = Diagonal(SVector{2^N,ComplexF64}(exp(factor * diagonal[i]) for i = 1:2^N))
+    apply_gate!(g_mat, state_vec, ts...)
+end
+function apply_gate!(
+    factor::ComplexF64,
+    state_vec::StateVector{T},
+    ts::Vararg{Int,N},
+) where {T<:Complex,N}
+    g_mat = Diagonal(SVector{2^N,ComplexF64}(factor for i = 1:2^N))
+    apply_gate!(g_mat, state_vec, ts...)
+end
 for (V, f) in ((true, :conj), (false, :identity))
     @eval begin
         apply_gate!(
@@ -77,77 +94,44 @@ for (V, f) in ((true, :conj), (false, :identity))
             t1::Int,
             t2::Int,
         ) where {T<:Complex} = apply_gate!(Val($V), ZZ(g.angle), state_vec, t1, t2)
-        function apply_gate!(
+        apply_gate!(
             ::Val{$V},
             g::MultiRZ,
             state_vec::StateVector{T},
             ts::Vararg{Int,N},
-        ) where {T<:Complex,N}
-            factor = -im * g.angle[1] / 2.0
-            r_mat = Braket.PauliEigenvalues(Val(N))
-            g_mat = Diagonal($f(SVector{2^N,ComplexF64}(exp(factor * r_mat[i]) for i = 1:2^N)))
-            apply_gate!(g_mat, state_vec, ts...)
-        end
-        function apply_gate!(
-            ::Val{$V},
-            g::MultiQubitPhaseShift{1},
-            state_vec::StateVector{T},
-            t::Int,
-        ) where {T<:Complex}
-            g_mat = $f(Diagonal(SVector{2, ComplexF64}(exp(im*g.angle[1]), exp(im*g.angle[1]))))
-            return apply_gate!(g_mat, state_vec, t)
-        end
-        function apply_gate!(
-            ::Val{$V},
-            g::MultiQubitPhaseShift{N},
-            state_vec::StateVector{T},
-            ts::Vararg{Int,N},
-        ) where {T<:Complex,N}
-            n_amps, endian_ts = get_amps_and_qubits(state_vec, ts...)
-            ordered_ts = sort(collect(endian_ts))
-            flip_list = map(0:2^N-1) do t
-                f_vals = Bool[(((1 << f_ix) & t) >> f_ix) for f_ix = 0:N-1]
-                return ordered_ts[f_vals]
-            end
-            factor = im * g.angle[1]
-            r_mat = ones(Float64, 2^N)
-            g_mat = Diagonal($f(SVector{2^N,ComplexF64}(exp(factor) * r_mat[i] for i = 1:2^N)))
-            apply_gate!(g_mat, state_vec, ts...)
-        end
+        ) where {T<:Complex,N} = apply_gate!($f(-im * g.angle[1] / 2.0), Braket.PauliEigenvalues(Val(N)), state_vec, ts...)
+        apply_gate!(::Val{$V}, g::MultiQubitPhaseShift{N}, state_vec::StateVector{T}, ts::Vararg{Int,N}) where {T<:Complex, N} = apply_gate!($f(exp(im*g.angle[1])), state_vec, ts...)
     end
 end
 
-for V in (false, true)
-    @eval begin
-        function apply_gate!(
-            ::Val{$V},
-            g::DoubleExcitation,
-            state_vec::StateVector{T},
-            t1::Int,
-            t2::Int,
-            t3::Int,
-            t4::Int,
-        ) where {T<:Complex}
-            n_amps, endian_ts = get_amps_and_qubits(state_vec, t1, t2, t3, t4)
-            ordered_ts = sort(collect(endian_ts))
-            cosϕ = cos(g.angle[1] / 2.0)
-            sinϕ = sin(g.angle[1] / 2.0)
-            e_t1, e_t2, e_t3, e_t4 = endian_ts
-            Threads.@threads for ix = 0:div(n_amps, 2^4)-1
-                padded_ix = pad_bits(ix, ordered_ts)
-                i0011 = flip_bits(padded_ix, (e_t3, e_t4)) + 1
-                i1100 = flip_bits(padded_ix, (e_t1, e_t2)) + 1
-                @inbounds begin
-                    amp0011 = state_vec[i0011]
-                    amp1100 = state_vec[i1100]
-                    state_vec[i0011] = cosϕ * amp0011 - sinϕ * amp1100
-                    state_vec[i1100] = sinϕ * amp0011 + cosϕ * amp1100
-                end
-            end
-            return
+function apply_gate!(
+    g::DoubleExcitation,
+    state_vec::StateVector{T},
+    t1::Int,
+    t2::Int,
+    t3::Int,
+    t4::Int,
+) where {T<:Complex}
+    n_amps, endian_ts = get_amps_and_qubits(state_vec, t1, t2, t3, t4)
+    ordered_ts = sort(collect(endian_ts))
+    cosϕ = cos(g.angle[1] / 2.0)
+    sinϕ = sin(g.angle[1] / 2.0)
+    e_t1, e_t2, e_t3, e_t4 = endian_ts
+    Threads.@threads for ix = 0:div(n_amps, 2^4)-1
+        padded_ix = pad_bits(ix, ordered_ts)
+        i0011 = flip_bits(padded_ix, (e_t3, e_t4)) + 1
+        i1100 = flip_bits(padded_ix, (e_t1, e_t2)) + 1
+        @inbounds begin
+            amp0011 = state_vec[i0011]
+            amp1100 = state_vec[i1100]
+            state_vec[i0011] = cosϕ * amp0011 - sinϕ * amp1100
+            state_vec[i1100] = sinϕ * amp0011 + cosϕ * amp1100
         end
     end
+    return
 end
+apply_gate!(::Val{true}, g::DoubleExcitation, state_vec::StateVector{T}, t1::Int, t2::Int, t3::Int, t4::Int) where {T<:Complex} = apply_gate!(g, state_vec, t1, t2, t3, t4)
+apply_gate!(::Val{false}, g::DoubleExcitation, state_vec::StateVector{T}, t1::Int, t2::Int, t3::Int, t4::Int) where {T<:Complex} = apply_gate!(g, state_vec, t1, t2, t3, t4)
 
 struct Control{G<:Gate, B} <: Gate
     g::G

src/dm_simulator.jl

@@ -39,10 +39,10 @@ mutable struct DensityMatrixSimulator{T,S} <:
     end
 end
 function init(
-    ::Type{DensityMatrixSimulator{T,S}},
+    t::Type{S},
     qubit_count::Int,
-) where {T,S<:AbstractMatrix{T}}
-    dm = S(undef, 2^qubit_count, 2^qubit_count)
+) where {T<:Complex,S<:AbstractMatrix{T}}
+    dm = t(undef, 2^qubit_count, 2^qubit_count)
     fill!(dm, zero(T))
     dm[1, 1] = one(T)
     return dm
@@ -51,7 +51,7 @@ function DensityMatrixSimulator{T,S}(
     qubit_count::Int,
     shots::Int,
 ) where {T,S<:AbstractDensityMatrix{T}}
-    dm = init(DensityMatrixSimulator{T,S}, qubit_count)
+    dm = init(S, qubit_count)
     return DensityMatrixSimulator{T,S}(dm, qubit_count, shots)
 end
 """
@@ -70,10 +70,12 @@ Braket.qubit_count(dms::DensityMatrixSimulator) = dms.qubit_count
 Query the properties and capabilities of a `DensityMatrixSimulator`, including which gates and result types are supported and the minimum and maximum shot and qubit counts.
 """
 Braket.properties(d::DensityMatrixSimulator) = dm_props
-supported_operations(d::DensityMatrixSimulator) =
-    dm_props.action["braket.ir.openqasm.program"].supportedOperations
-supported_result_types(d::DensityMatrixSimulator) =
-    dm_props.action["braket.ir.openqasm.program"].supportedResultTypes
+supported_operations(d::DensityMatrixSimulator, ::Val{:OpenQASM}) = dm_props.action["braket.ir.openqasm.program"].supportedOperations
+supported_operations(d::DensityMatrixSimulator, ::Val{:JAQCD}) = dm_props.action["braket.ir.jaqcd.program"].supportedOperations
+supported_operations(d::DensityMatrixSimulator) = supported_operations(d::DensityMatrixSimulator, Val(:OpenQASM))
+supported_result_types(d::DensityMatrixSimulator, ::Val{:OpenQASM}) = dm_props.action["braket.ir.openqasm.program"].supportedResultTypes
+supported_result_types(d::DensityMatrixSimulator, ::Val{:JAQCD}) = dm_props.action["braket.ir.jaqcd.program"].supportedResultTypes
+supported_result_types(d::DensityMatrixSimulator) = supported_result_types(d::DensityMatrixSimulator, Val(:OpenQASM))
 Braket.device_id(dms::DensityMatrixSimulator) = "braket_dm_v2"
 Braket.name(dms::DensityMatrixSimulator) = "DensityMatrixSimulator"
 Base.show(io::IO, dms::DensityMatrixSimulator) =
@@ -169,24 +171,14 @@ function evolve!(
     return dms
 end
 
-for (gate, obs) in (
-    (:X, :(Braket.Observables.X)),
-    (:Y, :(Braket.Observables.Y)),
-    (:Z, :(Braket.Observables.Z)),
-    (:I, :(Braket.Observables.I)),
-    (:H, :(Braket.Observables.H)),
-)
-    @eval begin
-        function apply_observable!(
-            observable::$obs,
-            dm::S,
-            targets,
-        ) where {T<:Complex,S<:AbstractDensityMatrix{T}}
-            reshaped_dm = reshape(dm, length(dm))
-            foreach(target->apply_gate!($gate(), reshaped_dm, target), targets)
-            return dm
-        end
-    end
+function apply_observable!(
+    gate::G,
+    dm::S,
+    targets,
+) where {T<:Complex,S<:AbstractDensityMatrix{T}, G<:Gate}
+    reshaped_dm = reshape(dm, length(dm))
+    foreach(target->apply_gate!(gate, reshaped_dm, target), targets)
+    return dm
 end
 function apply_observable!(
     observable::Braket.Observables.HermitianObservable,
@@ -233,8 +225,7 @@ end
 function apply_observables!(dms::DensityMatrixSimulator, observables)
     !isempty(dms._density_matrix_after_observables) &&
         error("observables have already been applied.")
-    diag_gates = [diagonalizing_gates(observable...) for observable in observables]
-    operations = reduce(vcat, diag_gates)
+    operations = mapreduce(obs->diagonalizing_gates(obs...), vcat, observables)
     dms._density_matrix_after_observables = deepcopy(dms.density_matrix)
     reshaped_dm = reshape(dms._density_matrix_after_observables, length(dms.density_matrix))
     for operation in operations