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.
It’s just one example

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  = local.rg_name
  storage_account_type = local.datadisk_type
  create_option        = "Empty"
  disk_size_gb         = local.disk_size

  tags = {
    environment = local.environment
  }
}

what incase if i have to dynamically change the size of disk per VM?