/** @file
  AML Tree Iterator.

  Copyright (c) 2019 - 2020, Arm Limited. All rights reserved.<BR>

  SPDX-License-Identifier: BSD-2-Clause-Patent
**/

#include <AmlNodeDefines.h>
#include <Tree/AmlTreeIterator.h>

#include <AmlCoreInterface.h>
#include <Tree/AmlTreeTraversal.h>

/** Iterator to traverse the tree.

  This is an internal structure.
*/
typedef struct AmlTreeInternalIterator {
  /// External iterator structure, containing the external APIs.
  /// Must be the first field.
  AML_TREE_ITERATOR         Iterator;

  // Note: The following members of this structure are opaque to the users
  //       of the Tree iterator APIs.

  /// Pointer to the node on which the iterator has been initialized.
  CONST  AML_NODE_HEADER    *InitialNode;

  /// Pointer to the current node.
  CONST  AML_NODE_HEADER    *CurrentNode;

  /// Iteration mode.
  /// Allow to choose how to traverse the tree/choose which node is next.
  EAML_ITERATOR_MODE        Mode;
} AML_TREE_ITERATOR_INTERNAL;

/** Get the current node of an iterator.

  @param  [in]  Iterator  Pointer to an iterator.
  @param  [out] OutNode   Pointer holding the current node.

  @retval EFI_SUCCESS             The function completed successfully.
  @retval EFI_INVALID_PARAMETER   Invalid parameter.
**/
STATIC
EFI_STATUS
EFIAPI
AmlIteratorGetNode (
  IN  AML_TREE_ITERATOR  *Iterator,
  OUT AML_NODE_HEADER    **OutNode
  )
{
  AML_TREE_ITERATOR_INTERNAL  *InternalIterator;

  InternalIterator = (AML_TREE_ITERATOR_INTERNAL *)Iterator;

  // CurrentNode can be NULL, but InitialNode cannot.
  if ((OutNode == NULL)                                       ||
      (InternalIterator == NULL)                              ||
      (InternalIterator->Mode <= EAmlIteratorUnknown)         ||
      (InternalIterator->Mode >= EAmlIteratorModeMax)         ||
      !IS_AML_NODE_VALID (InternalIterator->InitialNode)      ||
      ((InternalIterator->CurrentNode != NULL)                &&
       !IS_AML_NODE_VALID (InternalIterator->CurrentNode)))
  {
    ASSERT (0);
    return EFI_INVALID_PARAMETER;
  }

  *OutNode = (AML_NODE_HEADER *)InternalIterator->CurrentNode;

  return EFI_SUCCESS;
}

/** Move the current node of the iterator to the next node,
    according to the iteration mode selected.

  If NextNode is not NULL, return the next node.

  @param  [in]  Iterator    Pointer to an iterator.
  @param  [out] NextNode    If not NULL, updated to the next node.

  @retval EFI_SUCCESS             The function completed successfully.
  @retval EFI_INVALID_PARAMETER   Invalid parameter.
**/
STATIC
EFI_STATUS
EFIAPI
AmlIteratorGetNextLinear (
  IN  AML_TREE_ITERATOR  *Iterator,
  OUT AML_NODE_HEADER    **NextNode
  )
{
  AML_TREE_ITERATOR_INTERNAL  *InternalIterator;

  InternalIterator = (AML_TREE_ITERATOR_INTERNAL *)Iterator;

  // CurrentNode can be NULL, but InitialNode cannot.
  if ((InternalIterator == NULL)                              ||
      (InternalIterator->Mode != EAmlIteratorLinear)          ||
      !IS_AML_NODE_VALID (InternalIterator->InitialNode)      ||
      !IS_AML_NODE_VALID (InternalIterator->CurrentNode))
  {
    ASSERT (0);
    return EFI_INVALID_PARAMETER;
  }

  // Get the next node according to the iteration mode.
  InternalIterator->CurrentNode = AmlGetNextNode (
                                    InternalIterator->CurrentNode
                                    );

  if (NextNode != NULL) {
    *NextNode = (AML_NODE_HEADER *)InternalIterator->CurrentNode;
  }

  return EFI_SUCCESS;
}

/** Move the current node of the iterator to the previous node,
    according to the iteration mode selected.

  If PrevNode is not NULL, return the previous node.

  @param  [in]  Iterator    Pointer to an iterator.
  @param  [out] PrevNode    If not NULL, updated to the previous node.

  @retval EFI_SUCCESS             The function completed successfully.
  @retval EFI_INVALID_PARAMETER   Invalid parameter.
**/
STATIC
EFI_STATUS
EFIAPI
AmlIteratorGetPreviousLinear (
  IN  AML_TREE_ITERATOR  *Iterator,
  OUT AML_NODE_HEADER    **PrevNode
  )
{
  AML_TREE_ITERATOR_INTERNAL  *InternalIterator;

  InternalIterator = (AML_TREE_ITERATOR_INTERNAL *)Iterator;

  // CurrentNode can be NULL, but InitialNode cannot.
  if ((InternalIterator == NULL)                              ||
      (InternalIterator->Mode != EAmlIteratorLinear)          ||
      !IS_AML_NODE_VALID (InternalIterator->InitialNode)      ||
      !IS_AML_NODE_VALID (InternalIterator->CurrentNode))
  {
    ASSERT (0);
    return EFI_INVALID_PARAMETER;
  }

  // Get the previous node according to the iteration mode.
  InternalIterator->CurrentNode = AmlGetPreviousNode (
                                    InternalIterator->CurrentNode
                                    );
  if (PrevNode != NULL) {
    *PrevNode = (AML_NODE_HEADER *)InternalIterator->CurrentNode;
  }

  return EFI_SUCCESS;
}

/** Move the current node of the iterator to the next node,
    according to the iteration mode selected.

  If NextNode is not NULL, return the next node.

  @param  [in]  Iterator    Pointer to an iterator.
  @param  [out] NextNode    If not NULL, updated to the next node.

  @retval EFI_SUCCESS             The function completed successfully.
  @retval EFI_INVALID_PARAMETER   Invalid parameter.
**/
STATIC
EFI_STATUS
EFIAPI
AmlIteratorGetNextBranch (
  IN  AML_TREE_ITERATOR  *Iterator,
  OUT AML_NODE_HEADER    **NextNode
  )
{
  AML_TREE_ITERATOR_INTERNAL  *InternalIterator;
  AML_NODE_HEADER             *Node;

  InternalIterator = (AML_TREE_ITERATOR_INTERNAL *)Iterator;

  // CurrentNode can be NULL, but InitialNode cannot.
  if ((InternalIterator == NULL)                              ||
      (InternalIterator->Mode != EAmlIteratorBranch)          ||
      !IS_AML_NODE_VALID (InternalIterator->InitialNode)      ||
      !IS_AML_NODE_VALID (InternalIterator->CurrentNode))
  {
    ASSERT (0);
    return EFI_INVALID_PARAMETER;
  }

  Node = AmlGetNextNode (InternalIterator->CurrentNode);
  // Check whether NextNode is a sibling of InitialNode.
  if (AmlGetParent (Node) ==
      AmlGetParent ((AML_NODE_HEADER *)InternalIterator->InitialNode))
  {
    Node = NULL;
  }

  InternalIterator->CurrentNode = Node;

  if (NextNode != NULL) {
    *NextNode = Node;
  }

  return EFI_SUCCESS;
}

/** Move the current node of the iterator to the previous node,
    according to the iteration mode selected.

  If PrevNode is not NULL, return the previous node.

  @param  [in]  Iterator    Pointer to an iterator.
  @param  [out] PrevNode    If not NULL, updated to the previous node.

  @retval EFI_SUCCESS             The function completed successfully.
  @retval EFI_INVALID_PARAMETER   Invalid parameter.
**/
STATIC
EFI_STATUS
EFIAPI
AmlIteratorGetPreviousBranch (
  IN  AML_TREE_ITERATOR  *Iterator,
  OUT AML_NODE_HEADER    **PrevNode
  )
{
  AML_TREE_ITERATOR_INTERNAL  *InternalIterator;
  AML_NODE_HEADER             *Node;

  InternalIterator = (AML_TREE_ITERATOR_INTERNAL *)Iterator;

  // CurrentNode can be NULL, but InitialNode cannot.
  if ((InternalIterator == NULL)                              ||
      (InternalIterator->Mode != EAmlIteratorBranch)          ||
      !IS_AML_NODE_VALID (InternalIterator->InitialNode)      ||
      !IS_AML_NODE_VALID (InternalIterator->CurrentNode))
  {
    ASSERT (0);
    return EFI_INVALID_PARAMETER;
  }

  Node = AmlGetPreviousNode (InternalIterator->CurrentNode);
  // Check whether PreviousNode is a sibling of InitialNode.
  if (AmlGetParent (Node) ==
      AmlGetParent ((AML_NODE_HEADER *)InternalIterator->InitialNode))
  {
    Node = NULL;
  }

  InternalIterator->CurrentNode = Node;

  if (PrevNode != NULL) {
    *PrevNode = Node;
  }

  return EFI_SUCCESS;
}

/** Initialize an iterator.

  Note: The caller must call AmlDeleteIterator () to free the memory
        allocated for the iterator.

  @param  [in]  Node          Pointer to the node.
  @param  [in]  IteratorMode  Selected mode to traverse the tree.
  @param  [out] IteratorPtr   Pointer holding the created iterator.

  @retval EFI_SUCCESS             The function completed successfully.
  @retval EFI_INVALID_PARAMETER   Invalid parameter.
  @retval EFI_OUT_OF_RESOURCES    Could not allocate memory.
**/
EFI_STATUS
EFIAPI
AmlInitializeIterator (
  IN   AML_NODE_HEADER     *Node,
  IN   EAML_ITERATOR_MODE  IteratorMode,
  OUT  AML_TREE_ITERATOR   **IteratorPtr
  )
{
  AML_TREE_ITERATOR_INTERNAL  *InternalIterator;

  if (!IS_AML_NODE_VALID (Node)             ||
      (IteratorMode <= EAmlIteratorUnknown) ||
      (IteratorMode >= EAmlIteratorModeMax) ||
      (IteratorPtr == NULL))
  {
    ASSERT (0);
    return EFI_INVALID_PARAMETER;
  }

  *IteratorPtr     = NULL;
  InternalIterator = (AML_TREE_ITERATOR_INTERNAL *)AllocateZeroPool (
                                                     sizeof (
                                                             AML_TREE_ITERATOR_INTERNAL
                                                             )
                                                     );
  if (InternalIterator == NULL) {
    ASSERT (0);
    return EFI_OUT_OF_RESOURCES;
  }

  InternalIterator->InitialNode      = Node;
  InternalIterator->CurrentNode      = Node;
  InternalIterator->Mode             = IteratorMode;
  InternalIterator->Iterator.GetNode = AmlIteratorGetNode;

  switch (InternalIterator->Mode) {
    case EAmlIteratorLinear:
    {
      InternalIterator->Iterator.GetNext     = AmlIteratorGetNextLinear;
      InternalIterator->Iterator.GetPrevious = AmlIteratorGetPreviousLinear;
      break;
    }
    case EAmlIteratorBranch:
    {
      InternalIterator->Iterator.GetNext     = AmlIteratorGetNextBranch;
      InternalIterator->Iterator.GetPrevious = AmlIteratorGetPreviousBranch;
      break;
    }
    default:
    {
      ASSERT (0);
      FreePool (InternalIterator);
      return EFI_INVALID_PARAMETER;
    }
  } // switch

  *IteratorPtr = &InternalIterator->Iterator;

  return EFI_SUCCESS;
}

/** Delete an iterator.

  Note: The caller must have first initialized the iterator with the
        AmlInitializeIterator () function.

  @param  [in]  Iterator  Pointer to an iterator.

  @retval EFI_SUCCESS             The function completed successfully.
  @retval EFI_INVALID_PARAMETER   Invalid parameter.
**/
EFI_STATUS
EFIAPI
AmlDeleteIterator (
  IN  AML_TREE_ITERATOR  *Iterator
  )
{
  if (Iterator == NULL) {
    ASSERT (0);
    return EFI_INVALID_PARAMETER;
  }

  FreePool (Iterator);

  return EFI_SUCCESS;
}