diff --git a/benchmark/prof.jl b/benchmark/prof.jl
index 1eb99ab..74f1e39 100644
--- a/benchmark/prof.jl
+++ b/benchmark/prof.jl
@@ -9,8 +9,8 @@ using Profile
 using PProf
 using JET
 
-include("../test/problems/goddard.jl")
-#include("../test/problems/double_integrator.jl")
+#include("../test/problems/goddard.jl")
+include("../test/problems/simple_integrator.jl")
 
 # local version of mayer cost
 function local_mayer(obj, x0, xf, v)
@@ -20,8 +20,8 @@ end
 
 function init(;grid_size, disc_method)
     #prob = goddard_all()
-    prob = goddard()
-    #prob = double_integrator_mintf()
+    #prob = goddard()
+    prob = simple_integrator()
     ocp = prob[:ocp]
     docp = CTDirect.DOCP(ocp, grid_size=grid_size, time_grid=CTDirect.__time_grid(), disc_method=string(disc_method))
     xu = CTDirect.DOCP_initial_guess(docp)
diff --git a/src/midpoint.jl b/src/midpoint.jl
index 330a245..e6b63d3 100644
--- a/src/midpoint.jl
+++ b/src/midpoint.jl
@@ -4,20 +4,20 @@ Internal layout for NLP variables:
 with the convention u([t_i,t_i+1[) = U_i and u(tf) = U_N-1
 =#
 
-# +++ TODO: use args
-# NB. could be defined as a generic IRK
+# NB. could be defined as a generic IRK in future IRK.jl for comparison purposes
+#butcher_a::Matrix{Float64}
+#butcher_b::Vector{Float64}
+#butcher_c::Vector{Float64}
+#return new(1, 0, hcat(0.5), [1], [0.5], "Implicit Midpoint aka Gauss-Legendre collocation for s=1, 2nd order, symplectic")
 struct Midpoint <: Discretization
 
     stage::Int
-    additional_controls::Int
-    butcher_a::Matrix{Float64}
-    butcher_b::Vector{Float64}
-    butcher_c::Vector{Float64}
+    additional_controls::Int  # add control at tf
     info::String
 
-    # constructor    
-    function Midpoint(dim_NLP_x, dim_NLP_u, dim_NLP_steps) 
-        return new(1, 0, hcat(0.5), [1], [0.5], "Implicit Midpoint aka Gauss-Legendre collocation for s=1, 2nd order, symplectic")
+    # constructor
+    function Midpoint()
+        return new(1, 0, "Implicit Midpoint aka Gauss-Legendre collocation for s=1, 2nd order, symplectic")
     end
 end
 
@@ -61,48 +61,70 @@ end
 
 
 function get_stagevars_at_time_step(xu, docp::DOCP{Midpoint}, i)
-    nx = docp.dim_NLP_x
-    m = docp.dim_NLP_u
-    N = docp.dim_NLP_steps
-    offset = (nx*(1+docp.discretization.stage) + m) * (i-1)
+    offset = (i-1) * (docp.dim_NLP_x*(1+docp.discretization.stage) + docp.dim_NLP_u) + docp.dim_NLP_x + docp.dim_NLP_u
     # retrieve vector stage variable (except at final time)
-    if i < N+1
-        return @view xu[(offset + nx + m + 1):(offset + nx + m + nx) ]
+    if i < docp.dim_NLP_steps+1
+        return @view xu[(offset + 1):(offset + docp.dim_NLP_x)]
     else
-        return nothing
+        return @view xu[1:docp.dim_NLP_x] # unused but keep same type !
     end
 end
 
+
 function set_state_at_time_step!(xu, x_init, docp::DOCP{Midpoint}, i)
-    nx = docp.dim_NLP_x
-    n = docp.dim_OCP_x
-    m = docp.dim_NLP_u
-    offset = (nx*(1+docp.discretization.stage) + m) * (i-1)
+    offset = (i-1) * (docp.dim_NLP_x*(1+docp.discretization.stage) + docp.dim_NLP_u)
     # initialize only actual state variables from OCP (not lagrange state)
     if !isnothing(x_init)
-        xu[(offset + 1):(offset + n)] .= x_init
+        xu[(offset + 1):(offset + docp.dim_OCP_x)] .= x_init
     end
 end
 
+
 function set_control_at_time_step!(xu, u_init, docp::DOCP{Midpoint}, i)
-    nx = docp.dim_NLP_x
-    m = docp.dim_NLP_u
-    N = docp.dim_NLP_steps
-    if i < N+1
-        offset = (nx*(1+docp.discretization.stage) + m) * (i-1)
+    if i < docp.dim_NLP_steps+1
+        offset = (i-1) * (docp.dim_NLP_x*(1+docp.discretization.stage) + docp.dim_NLP_u) + docp.dim_NLP_x
     else
-        offset = (nx*(1+docp.discretization.stage) + m) * (i-2)
+        offset = (i-2) * (docp.dim_NLP_x*(1+docp.discretization.stage) + docp.dim_NLP_u) + docp.dim_NLP_x
     end
     if !isnothing(u_init)
-        xu[(offset + nx + 1):(offset + nx + m)] .= u_init
+        xu[(offset + 1):(offset + docp.dim_NLP_u)] .= u_init
     end
 end
 
 
 function setWorkArray(docp::DOCP{Midpoint}, xu, time_grid, v)
-    return nothing
+    # use work array to store all dynamics + lagrange costs
+    work = similar(xu, docp.dim_NLP_x * (docp.dim_NLP_steps))
+    if docp.dim_OCP_x > 1
+        xs = similar(xu, docp.dim_OCP_x)
+    end
+
+    # loop over time steps
+    for i = 1:docp.dim_NLP_steps
+        offset = (i-1) * docp.dim_NLP_x
+        ti = time_grid[i]
+        xi = get_OCP_state_at_time_step(xu, docp, i)
+        ui = get_OCP_control_at_time_step(xu, docp, i)
+        tip1 = time_grid[i+1]
+        xip1 = get_OCP_state_at_time_step(xu, docp, i+1)
+        ts = 0.5 * (ti + tip1)
+        # add a new getter ?
+        if docp.dim_OCP_x == 1
+            xs = 0.5 * (xi + xip1)
+        else
+            @. xs = 0.5 * (xi + xip1) # not ok for dim 1...
+        end
+        # OCP dynamics
+        docp.ocp.dynamics((@view work[offset+1:offset+docp.dim_OCP_x]), ts, xs, ui, v)
+        # lagrange cost
+        if docp.is_lagrange
+            docp.ocp.lagrange((@view work[offset+docp.dim_NLP_x:offset+docp.dim_NLP_x]), ts, xs, ui, v)
+        end
+    end
+    return work
 end
 
+
 """
 $(TYPEDSIGNATURES)
 
@@ -111,15 +133,13 @@ Convention: 1 <= i <= dim_NLP_steps (+1)
 """
 function setConstraintBlock!(docp::DOCP{Midpoint}, c, xu, v, time_grid, i, work)
 
-    disc = docp.discretization
-
     # offset for previous steps
     offset = (i-1)*(docp.dim_NLP_x * (1+docp.discretization.stage) + docp.dim_path_cons)
 
-    # variables
+    # 0. variables
     ti = time_grid[i]
     xi = get_OCP_state_at_time_step(xu, docp, i)
-    ui = get_OCP_control_at_time_step(xu, docp, i) 
+    ui = get_OCP_control_at_time_step(xu, docp, i)
 
     # 1. state equation
     if i <= docp.dim_NLP_steps
@@ -128,42 +148,33 @@ function setConstraintBlock!(docp::DOCP{Midpoint}, c, xu, v, time_grid, i, work)
         tip1 = time_grid[i+1]
         xip1 = get_OCP_state_at_time_step(xu, docp, i+1)
         hi = tip1 - ti
+        offset_dyn_i = (i-1)*docp.dim_NLP_x
 
         # midpoint rule
-        h_sum_bk = hi * disc.butcher_b[1] * ki
-        if docp.dim_OCP_x == 1
-            c[offset+1] = xip1 - (xi + h_sum_bk[1])
-        else
-            c[offset+1:offset+docp.dim_OCP_x] = xip1 - (xi + h_sum_bk[1:docp.dim_OCP_x])
-        end
+        @views @. c[offset+1:offset+docp.dim_OCP_x] = xip1 - (xi + hi * ki[1:docp.dim_OCP_x])
+
         if docp.is_lagrange
-            c[offset+docp.dim_NLP_x] = get_lagrange_state_at_time_step(xu, docp, i+1) - (get_lagrange_state_at_time_step(xu, docp, i) + h_sum_bk[docp.dim_NLP_x])
+            c[offset+docp.dim_NLP_x] = get_lagrange_state_at_time_step(xu, docp, i+1) - (get_lagrange_state_at_time_step(xu, docp, i) + hi * ki[docp.dim_NLP_x])
         end
         offset += docp.dim_NLP_x
 
         # stage equation at mid-step
-        ts = ti + hi * disc.butcher_c[1]
-        #xs = xi + hi * (disc.butcher_a[1][1] * ki)
-        xs = 0.5 * (xi + xip1) #compare bench
-        docp.ocp.dynamics((@view c[offset+1:offset+docp.dim_OCP_x]), ts, xs, ui, v)
-        @views c[offset+1:offset+docp.dim_OCP_x] = -c[offset+1:offset+docp.dim_OCP_x] + ki[1:docp.dim_OCP_x]
+        @views @. c[offset+1:offset+docp.dim_OCP_x] = ki[1:docp.dim_OCP_x] - work[offset_dyn_i+1:offset_dyn_i+docp.dim_OCP_x]
         if docp.is_lagrange
-            docp.ocp.lagrange((@view c[offset+docp.dim_NLP_x:offset+docp.dim_NLP_x]), ts, xs, ui, v)
-            @views c[offset+docp.dim_NLP_x] = -c[offset+docp.dim_NLP_x] + ki[docp.dim_NLP_x]
+            c[offset+docp.dim_NLP_x] = ki[docp.dim_NLP_x] - work[offset_dyn_i+docp.dim_NLP_x]
         end
         offset += docp.dim_NLP_x
     end
 
     # 2. path constraints
-    # Notes on allocations:.= seems similar
     if docp.dim_u_cons > 0
-        docp.control_constraints[2]((@view c[offset+1:offset+docp.dim_u_cons]),ti, docp._u(ui), v)
+        docp.control_constraints[2]((@view c[offset+1:offset+docp.dim_u_cons]),ti, ui, v)
     end
     if docp.dim_x_cons > 0 
-        docp.state_constraints[2]((@view c[offset+docp.dim_u_cons+1:offset+docp.dim_u_cons+docp.dim_x_cons]),ti, docp._x(xi), v)
+        docp.state_constraints[2]((@view c[offset+docp.dim_u_cons+1:offset+docp.dim_u_cons+docp.dim_x_cons]),ti, xi, v)
     end
     if docp.dim_mixed_cons > 0 
-        docp.mixed_constraints[2]((@view c[offset+docp.dim_u_cons+docp.dim_x_cons+1:offset+docp.dim_u_cons+docp.dim_x_cons+docp.dim_mixed_cons]), ti, docp._x(xi), docp._u(ui), v)
+        docp.mixed_constraints[2]((@view c[offset+docp.dim_u_cons+docp.dim_x_cons+1:offset+docp.dim_u_cons+docp.dim_x_cons+docp.dim_mixed_cons]), ti, xi, ui, v)
     end
 
 end
diff --git a/src/problem.jl b/src/problem.jl
index b3d85ba..334b5c3 100644
--- a/src/problem.jl
+++ b/src/problem.jl
@@ -180,9 +180,9 @@ struct DOCP{T <: Discretization, X <: ScalVect, U <: ScalVect, V <: ScalVect}
 
         # parameter: discretization method
         if disc_method == "midpoint"
-            discretization = CTDirect.Midpoint(dim_NLP_x, dim_NLP_u, dim_NLP_steps)
+            discretization = CTDirect.Midpoint()
         elseif disc_method == "trapeze"
-            discretization = CTDirect.Trapeze(dim_NLP_x, dim_NLP_u)
+            discretization = CTDirect.Trapeze()
         else
             error("Unknown discretization method:", disc_method)
         end
diff --git a/src/trapeze.jl b/src/trapeze.jl
index faeb797..e45e707 100644
--- a/src/trapeze.jl
+++ b/src/trapeze.jl
@@ -11,7 +11,7 @@ struct Trapeze <: Discretization
     info::String
 
     # constructor
-    function Trapeze(dim_NLP_x, dim_NLP_u)
+    function Trapeze()
         return new(0, 1, "Implicit Trapeze aka Crank-Nicolson, 2nd order, A-stable")
     end
 end
@@ -113,15 +113,15 @@ function setConstraintBlock!(docp::DOCP{Trapeze}, c, xu, v, time_grid, i, work)
         # more variables
         tip1 = time_grid[i+1]
         xip1 = get_OCP_state_at_time_step(xu, docp, i+1)
+        half_hi = 0.5 * (tip1 - ti)
         offset_dyn_i = (i-1)*docp.dim_NLP_x
         offset_dyn_ip1 = i*docp.dim_NLP_x
 
         # trapeze rule (no allocations ^^)
-        halfstep = 0.5 * (tip1 - ti)
-        @views @. c[offset+1:offset+docp.dim_OCP_x] = xip1 - (xi + halfstep * (work[offset_dyn_i+1:offset_dyn_i+docp.dim_OCP_x] + work[offset_dyn_ip1+1:offset_dyn_ip1+docp.dim_OCP_x]))
+        @views @. c[offset+1:offset+docp.dim_OCP_x] = xip1 - (xi + half_hi * (work[offset_dyn_i+1:offset_dyn_i+docp.dim_OCP_x] + work[offset_dyn_ip1+1:offset_dyn_ip1+docp.dim_OCP_x]))
 
         if docp.is_lagrange
-            c[offset+docp.dim_NLP_x] = get_lagrange_state_at_time_step(xu, docp, i+1) - (get_lagrange_state_at_time_step(xu, docp, i) + 0.5 * (tip1 - ti) * (work[offset_dyn_i+docp.dim_NLP_x] + work[offset_dyn_ip1+docp.dim_NLP_x]))
+            c[offset+docp.dim_NLP_x] = get_lagrange_state_at_time_step(xu, docp, i+1) - (get_lagrange_state_at_time_step(xu, docp, i) + half_hi * (work[offset_dyn_i+docp.dim_NLP_x] + work[offset_dyn_ip1+docp.dim_NLP_x]))
         end
         offset += docp.dim_NLP_x
     end