Azure - Terraform dynamic disk counts index issue

We created an Azure VM module which can able to provision VM’s with multiple disks. In this scenario the disk count and the VM counts are dynamic.

resource "azurerm_managed_disk" "disk" {
 count = "${var.disk_count * var.vm_count}"
 name = "${element(azurerm_virtual_machine.instance.*.name, count.index % var.vm_count)}-disk${format("%02d", count.index / var.vm_count + 1)}"
 location = "${var.location}"
 resource_group_name = "${var.rg_name}"
 storage_account_type = "StandardSSD_LRS"
 create_option = "Empty"
 disk_size_gb = "${var.disk_size}"
}

With this, the initial build is working perfectly, but whenever we tried to scale it, the disk count index is getting out of order and trying to recreate the disks. Is there any better way to do this?

You can use for_each in resource.
You can possibly do better, but you can use this example of how a base.

locals {
  instances      = 3 # number of machies to be create (for use in name of disks)
  datadisk_count = 2 # number of datadisk to be created per machine
  location       = "West Europe"
  rg_name        = "test_disk"
  datadisk_type  = "Standard_LRS"
  disk_size      = 10
  name           = "nmv"
  environment    = "test"

  test_mv            = formatlist("%s%03d", local.name, range(1, (local.instances + 1), 1))
  test_machines_disk = { for k in local.test_mv : k => local.datadisk_count }

  final = flatten([
    for mv_name, count in local.test_machines_disk : [
      for i in range(count) : {
        index = format("datadisk%s_%s", (i + 1), mv_name)
      }
    ]
  ])
}


resource "azurerm_resource_group" "test" {
  name     = local.rg_name
  location = local.location
}

resource "azurerm_managed_disk" "datadisk" {
  for_each             = toset([for j in local.final : j.index])
  name                 = each.key
  location             = local.location
  resource_group_name  = azurerm_resource_group.test.name
  storage_account_type = local.datadisk_type
  create_option        = "Empty"
  disk_size_gb         = local.disk_size

  tags = {
    environment = local.environment
  }
}