compiling matrix row assignment

I have this relatively simple code, generated with another Mathematica code (this is just an example): (edit: see below for simpler version)

code = Compile[{{s, _Real, 1}},
Block[{d2s = {{4., -8., 4.}, {4., -8., 4.}, {4., -8., 4.}}.s,
coefficients, MblockCoefficients, Mblocks, M, Mbc, vParts, v},
coefficients = {{{-4.62924, -3.38288, -3.} s^2, {0., 0., 0.}, {1.,
1., 1.}, {-1., -0.5, 0.} + {-1.54308, -1.12763, -1.} s^3 +
d2s}};
vParts = -Last[Transpose[coefficients]];
v = Flatten[vParts];
MblockCoefficients = (Partition[#1, 2 + 1] &) /@Transpose[Most[Transpose[coefficients]]];
{{{1., 0., 0.}, {0., 1., 0.}, {0.,0., 1.}}, {{3., -4., 1.}, {1., 0., -1.}, {-1.,4., -3.}}, {{4., -8., 4.}, {4., -8., 4.}, {4., -8.,4.}}}}]] &,
MblockCoefficients, {2}];
M = ArrayFlatten[Mblocks];
v[[{3, 1}]] = {1. - s[[-1]], -0.2 - s[[1]]^3};
Mbc = {{0., 0., 1.}, {0. + 3. s[[1]]^2, 0., 0.}};
M[[{3, 1}]] = Mbc;
Partition[LinearSolve[M, v], 2 + 1]
]];


It takes a vector and produces another vector, and as you see it is compiled.

I want to minimise the calls to MainEvaluate. There are 4 in this case*, of which I understand 3: DiagonalMatrix, ArrayFlatten and LinearSolve cannot be compiled.

My question is about the 4th one, which comes from the line M[[{3,1}]] = Mbc. I thought that maybe this type of assignment on rows of a matrix cannot be compiled, but the following code:

code2 = Compile[{{m, _Real, 2}, {v, _Real, 2}},
Block[{mnew = m},
mnew[[{3,1}]] = v;
mnew
]];


compiles just fine, without any call to MainEvaluate.

So what is the problem in the first case, and is there a simple fix? This is in Mathematica 11.0.0.

EDIT: This I think is a minimal working example of the above code,

codeNotWorking=Compile[{{s, _Real, 1}}, Module[{M, Mbc},
M = ArrayFlatten[{{{{1., 0.}, {0., 1.}}, {{0., 0.}, {0., 0.}}},
{{{0., 0.}, {0., 0.}}, {{1., 0.}, {0.,1.}}}}];
Mbc = {{0., 0., 1., 0.}, {0. + 3. s[[1]]^2, 0., 0., 0.}};
M[[{3, 1}]] = Mbc;
], {{M, _Real, 2}, {Mbc, _Real, 2}}];


Again the same line doesn't compile. Instead if I first evaluate the ArrayFlatten and paste the result in the above code, it does compile. So the problem seems to be that because ArrayFlatten doesn't compile, the compiler doesn't know the dimensions of M. But I thought the third argument that I gave should give that information?

*I find these with the code

Needs["CompiledFunctionTools"]
printcode[comp_] := StringReplace[CompiledFunctionToolsCompilePrint[comp],
"MainEvaluate" -> Style["MainEvaluate", Red]]

• Maybe you meant to write M[[3, 1]] = Mbc; instead of M[[{3, 1}]] = Mbc;? That removes the fourth call to MainEvaluate for me... But no, that does not make sense. – Henrik Schumacher Sep 25 '18 at 16:18
• @HenrikSchumacher Exactly M[[3,1]] is something else and not what I want. I reduced it further and the issue seems to be that the compiler doesn't know the dimensions of M, even when I give them? – Jansen Sep 26 '18 at 13:27