core_update_state_impl Subroutine

  subroutine core_update_state_impl(plan, state, src_q, mpi)
    type(fmm_plan_type), intent(in) :: plan
    type(fmm_state_type), intent(inout) :: state
    real(dp), intent(in) :: src_q(:)
    type(mpi_context), intent(in), optional :: mpi
    integer(i32) :: n_target_nodes, source_idx, cell, rank, rank_count, block_begin, block_end
    integer(i32), allocatable :: block_counts(:), block_offsets(:)
    real(dp) :: total_charge, e_bottom, phi_gauge, z_gauge

    if (.not. plan%built) error stop 'FMM plan is not built.'
    if (size(src_q) /= plan%nsrc) error stop 'src_q size does not match plan.'

    state%ready = .false.
    call ensure_state_capacity(plan, state)
    if (plan%periodic_k0_ready) then
      total_charge = sum(src_q)
      z_gauge = plan%periodic_k0_plan%break_z(1)
      e_bottom = -0.5_dp*total_charge/(eps0*plan%periodic_k0_plan%area_xy)
      phi_gauge = 0.0_dp
      do source_idx = 1_i32, plan%nsrc
        phi_gauge = phi_gauge - 0.5_dp*src_q(source_idx)* &
                    (sum(plan%periodic_k0_plan%panel(source_idx)%z)/3.0_dp - z_gauge)/ &
                    (eps0*plan%periodic_k0_plan%area_xy)
      end do
      call refresh_periodic_zero_mode_state( &
        plan%periodic_k0_plan, src_q, e_bottom, z_gauge, phi_gauge, state%periodic_k0_state &
        )
      call refresh_periodic_shell_mean_state(plan%periodic_shell_mean, src_q, state%periodic_shell_mean_state)
    end if
    if (plan%periodic_root_operator_ready) then
      if (allocated(state%periodic_multipole)) then
        if (any(shape(state%periodic_multipole) /= [plan%ncoef, plan%periodic_root_source_count])) then
          deallocate (state%periodic_multipole)
        end if
      end if
      if (.not. allocated(state%periodic_multipole)) then
        allocate (state%periodic_multipole(plan%ncoef, plan%periodic_root_source_count))
      end if
      state%periodic_multipole = 0.0_dp
      do source_idx = 1_i32, plan%nsrc
        cell = plan%periodic_source_cell(source_idx)
        state%periodic_multipole(:, cell) = state%periodic_multipole(:, cell) + &
                                            src_q(source_idx)*plan%periodic_source_p2m_basis(:, source_idx)
      end do
      if (allocated(state%anchor_local)) then
        if (any(shape(state%anchor_local) /= [plan%ncoef, plan%periodic_operator_anchor_count])) then
          deallocate (state%anchor_local, state%anchor_local_active)
        end if
      end if
      if (.not. allocated(state%anchor_local)) then
        allocate (state%anchor_local(plan%ncoef, plan%periodic_operator_anchor_count), &
                  state%anchor_local_active(plan%periodic_operator_anchor_count))
      end if
      state%anchor_local = 0.0_dp
      state%anchor_local_active = 0_i32
    end if
    n_target_nodes = merge(plan%target_nnode, plan%nnode, plan%target_tree_ready)
    rank = 0_i32
    rank_count = 1_i32
    if (present(mpi)) call mpi_get_rank_size(rank, rank_count, mpi)
    call m2l_rank_blocks(plan, n_target_nodes, rank_count, block_counts, block_offsets)
    block_begin = block_offsets(rank + 1_i32) + 1_i32
    block_end = block_offsets(rank + 1_i32) + block_counts(rank + 1_i32)
    !$omp parallel default(none) shared(plan, state, src_q, n_target_nodes, block_begin, block_end)
    if (associated(state%multipole_active)) then
      !$omp workshare
      state%multipole_active = 0_i32
      !$omp end workshare
    end if
    if (associated(state%local_active)) then
      !$omp workshare
      state%local_active = 0_i32
      !$omp end workshare
    end if
    if (plan%nsrc > 0_i32) then
      !$omp workshare
      state%src_q = src_q
      !$omp end workshare
    end if
    if (plan%nsource_leaf <= 0_i32 .and. associated(state%multipole)) then
      !$omp workshare
      state%multipole = 0.0d0
      !$omp end workshare
    end if
    if (plan%m2l_pair_count <= 0_i32 .and. associated(state%local)) then
      !$omp workshare
      state%local = 0.0d0
      !$omp end workshare
    end if

    call p2m_leaf_moments(plan, state)
    call m2m_upward_pass(plan, state)
    if (n_target_nodes > 0_i32) call m2l_accumulate(plan, state, block_begin, block_end)
    !$omp end parallel

    if (rank_count > 1_i32 .and. n_target_nodes > 0_i32 .and. plan%m2l_pair_count > 0_i32) then
      call mpi_allgatherv_inplace_real_dp(mpi, state%local, block_counts*plan%ncoef, block_offsets*plan%ncoef)
      call mpi_allgatherv_inplace_i32(mpi, state%local_active, block_counts, block_offsets)
    end if

    !$omp parallel default(none) shared(plan, state, n_target_nodes)
    if (n_target_nodes > 0_i32) then
      call apply_periodic_target_correction(plan, state)
      call l2l_downward_pass(plan, state)
    end if
    !$omp end parallel

    state%ready = .true.
    state%update_count = state%update_count + 1_i32
  end subroutine core_update_state_impl