fun replaceToMock(
    assembleModel: UtModel,
    shouldMock: (ClassId) -> Boolean
): UtModel = with(assembleModel) {
    if (this !is UtAssembleModel) return this

    if (shouldMock(classId)) {
        UtCompositeModel(id, classId, isMock = true).apply {
            modificationsChain.forEach {
                if (it is UtDirectSetFieldModel) {
                    fields[it.fieldId] = replaceToMock(it.fieldModel, shouldMock)
                }

                if (it is UtExecutableCallModel && it.executable is FuzzerMockableMethodId) {
                    (it.executable as FuzzerMockableMethodId).mock().forEach { (executionId, models) ->
                        mocks[executionId] = models.map { p -> replaceToMock(p, shouldMock) }
                    }
                }
            }
        }
    } else {
        val models = modificationsChain.map { call ->
            var mockedStatementModel: UtStatementModel? = null

            if (call is UtDirectSetFieldModel) {
                val mock = replaceToMock(call.fieldModel, shouldMock)

                if (mock != call.fieldModel) {
                    mockedStatementModel = UtDirectSetFieldModel(call.instance, call.fieldId, mock)
                }
            } else if (call is UtExecutableCallModel) {
                val params = call.params.map { m -> replaceToMock(m, shouldMock) }

                if (params != call.params) {
                    mockedStatementModel = UtExecutableCallModel(call.instance, call.executable, params)
                }
            }

            mockedStatementModel ?: call
        }

        with(assembleModel) {
            UtAssembleModel(id, classId, modelName, instantiationCall, origin) { models }
        }
    }
}