////////////////////////////////
// Functions: Arity Node

// construction

MR4TH_SYMBOL DS_ArityNode*
ds_aritynode_empty(Arena *arena){
  DS_ArityNode *node = push_array(arena, DS_ArityNode, 1);
  node->kind = DS_ArityKind_Empty;
  return(node);
}

MR4TH_SYMBOL DS_ArityNode*
ds_aritynode_singleton(Arena *arena){
  DS_ArityNode *node = push_array(arena, DS_ArityNode, 1);
  node->kind = DS_ArityKind_Singleton;
  node->coefficient = 1;
  return(node);
}

MR4TH_SYMBOL DS_ArityNode*
ds_aritynode_arrow(Arena *arena, DS_ArityNode *domain, DS_ArityNode *codomain){
  DS_ArityNode *node = 0;
  if (domain->kind == DS_ArityKind_Empty ||
      codomain->kind == DS_ArityKind_Empty){
    node = codomain;
  }
  else{
    node = push_array(arena, DS_ArityNode, 1);
    node->kind = DS_ArityKind_Arrow;
    node->coefficient = 1;
    node->arrow.domain = domain;
    node->arrow.codomain = codomain;
  }
  return(node);
}

MR4TH_SYMBOL DS_ArityNode*
ds_aritynode_paste(Arena *arena, DS_ArityNode *a, DS_ArityNode *b){
  ArenaTemp scratch = arena_get_scratch(&arena, 1);
  DS_ArityNode *node = 0;
  
  if (a->kind == DS_ArityKind_Empty){
    node = b;
  }
  else if (b->kind == DS_ArityKind_Empty){
    node = a;
  }
  else{
    // count parts
    U64 max_part_count = 0;
    if (a->kind != DS_ArityKind_Tuple){
      max_part_count += 1;
    }
    else{
      max_part_count += a->tuple.part_count;
    }
    if (b->kind != DS_ArityKind_Tuple){
      max_part_count += 1;
    }
    else{
      max_part_count += b->tuple.part_count;
    }
    
    // fill parts
    DS_ArityNode **parts = push_array(scratch.arena, DS_ArityNode*, max_part_count);
    DS_ArityNode **part_ptr = parts;
    
    if (a->kind != DS_ArityKind_Tuple){
      *part_ptr = ds_aritynode__part_copy(arena, a);
      (*part_ptr)->coefficient = 1;
      part_ptr += 1;
    }
    else{
      for (U64 i = 0; i < a->tuple.part_count; i += 1){
        DS_ArityNode *a_part = a->tuple.parts[i];
        *part_ptr = ds_aritynode__part_copy(arena, a_part);
        (*part_ptr)->coefficient = a_part->coefficient;
        part_ptr += 1;
      }
    }
    
    if (b->kind != DS_ArityKind_Tuple){
      DS_ArityNode *match_pos = ds_aritynode__part_find_match(parts, part_ptr, b);
      if (match_pos != 0){
        match_pos->coefficient += 1;
      }
      else{
        *part_ptr = ds_aritynode__part_copy(arena, b);
        (*part_ptr)->coefficient = 1;
        part_ptr += 1;
      }
    }
    else{
      for (U64 j = 0; j < b->tuple.part_count; j += 1){
        DS_ArityNode *b_part = b->tuple.parts[j];
        DS_ArityNode *match_pos = ds_aritynode__part_find_match(parts, part_ptr, b_part);
        if (match_pos != 0){
          match_pos->coefficient += b_part->coefficient;
        }
        else{
          *part_ptr = ds_aritynode__part_copy(arena, b_part);
          (*part_ptr)->coefficient = b_part->coefficient;
          part_ptr += 1;
        }
      }
    }
    
    // final part count
    U64 part_count = (U64)(part_ptr - parts);
    
    // fill result node
    Assert(part_count > 1 ||
           part_count == 1 && parts[0]->coefficient > 1);
    
    DS_ArityNode **parts_final = push_array(arena, DS_ArityNode*, part_count);
    MemoryCopy(parts_final, parts, sizeof(*parts)*part_count);
    
    node = push_array(arena, DS_ArityNode, 1);
    node->kind = DS_ArityKind_Tuple;
    node->coefficient = 1;
    node->tuple.parts = parts_final;
    node->tuple.part_count = part_count;
  }
  
  arena_release_scratch(&scratch);
  
  return(node);
}

MR4TH_SYMBOL DS_ArityNode*
ds_aritynode_sum(Arena *arena, ...){
  va_list args;
  va_start(args, arena);
  
  va_list args2;
  va_copy(args2, args);
  
  ArenaTemp scratch = arena_get_scratch(&arena, 1);
  
  // count parts
  U64 max_part_count = 0;
  for (;;){
    U32 multiplier = va_arg(args, U32);
    if (multiplier == 0){
      break;
    }
    DS_ArityNode *arity = va_arg(args, DS_ArityNode*);
    if (arity->kind == DS_ArityKind_Empty){
      // do nothing
    }
    else if (arity->kind != DS_ArityKind_Tuple){
      max_part_count += 1;
    }
    else{
      max_part_count += arity->tuple.part_count;
    }
  }
  
  // fill parts
  DS_ArityNode **parts = push_array(scratch.arena, DS_ArityNode*, max_part_count);
  DS_ArityNode **part_ptr = parts;
  
  for (;;){
    U32 multiplier = va_arg(args2, U32);
    if (multiplier == 0){
      break;
    }
    DS_ArityNode *arity = va_arg(args2, DS_ArityNode*);
    
    if (arity->kind == DS_ArityKind_Empty){
      // do nothing
    }
    else if (arity->kind != DS_ArityKind_Tuple){
      DS_ArityNode *match_pos = ds_aritynode__part_find_match(parts, part_ptr, arity);
      if (match_pos != 0){
        match_pos->coefficient += multiplier;
      }
      else{
        *part_ptr = ds_aritynode__part_copy(arena, arity);
        (*part_ptr)->coefficient = multiplier;
        part_ptr += 1;
      }
    }
    else{
      for (U64 j = 0; j < arity->tuple.part_count; j += 1){
        DS_ArityNode *b_part = arity->tuple.parts[j];
        DS_ArityNode *match_pos = ds_aritynode__part_find_match(parts, part_ptr, b_part);
        if (match_pos != 0){
          match_pos->coefficient += multiplier*b_part->coefficient;
        }
        else{
          *part_ptr = ds_aritynode__part_copy(arena, b_part);
          (*part_ptr)->coefficient = multiplier*b_part->coefficient;
          part_ptr += 1;
        }
      }
    }
  }
  
  // final part count
  U64 part_count = (U64)(part_ptr - parts);
  
  // fill result node
  DS_ArityNode *node = 0;
  if (part_count == 0){
    node = ds_aritynode_empty(arena);
  }
  else if (part_count == 1 &&
           parts[0]->coefficient == 1){
    node = parts[0];
  }
  else{
    DS_ArityNode **parts_final = push_array(arena, DS_ArityNode*, part_count);
    MemoryCopy(parts_final, parts, sizeof(*parts)*part_count);
    
    node = push_array(arena, DS_ArityNode, 1);
    node->kind = DS_ArityKind_Tuple;
    node->coefficient = 1;
    node->tuple.parts = parts_final;
    node->tuple.part_count = part_count;
  }
  
  arena_release_scratch(&scratch);
  
  va_end(args2);
  va_end(args);
  
  return(node);
}

MR4TH_SYMBOL DS_ArityNode*
ds_aritynode_ktuple(Arena *arena, U32 coefficient){
  DS_ArityNode *result = 0;
  
  if (coefficient > 1){
    result = push_array(arena, DS_ArityNode, 1);
    DS_ArityNode **parts = push_array(arena, DS_ArityNode*, 1);
    DS_ArityNode *singleton = push_array(arena, DS_ArityNode, 1);
    
    result->kind = DS_ArityKind_Tuple;
    result->coefficient = 1;
    result->tuple.parts = parts;
    result->tuple.part_count = 1;
    
    parts[0] = singleton;
    
    singleton->kind = DS_ArityKind_Singleton;
    singleton->coefficient = coefficient;
  }
  else if (coefficient == 1){
    result = ds_aritynode_singleton(arena);
  }
  else{
    result = ds_aritynode_empty(arena);
  }
  
  return(result);
}

// operations

MR4TH_SYMBOL DS_ArityNode*
ds_aritynode_copy(Arena *arena, DS_ArityNode *arity){
  DS_ArityNode *result = push_array(arena, DS_ArityNode, 1);
  result->kind = arity->kind;
  result->coefficient = arity->coefficient;
  switch (result->kind){
    case DS_ArityKind_Arrow:
    {
      result->arrow.domain = ds_aritynode_copy(arena, arity->arrow.domain);
      result->arrow.codomain = ds_aritynode_copy(arena, arity->arrow.codomain);
    }break;
    
    case DS_ArityKind_Tuple:
    {
      U64 part_count = arity->tuple.part_count;
      DS_ArityNode **parts = push_array(arena, DS_ArityNode*, part_count);
      for (U64 i = 0; i < part_count; i += 1){
        parts[i] = ds_aritynode_copy(arena, arity->tuple.parts[i]);
      }
      result->tuple.parts = parts;
      result->tuple.part_count = part_count;
    }break;
  }
  return(result);
}

MR4TH_SYMBOL B32
ds_aritynode_equal(DS_ArityNode *a, DS_ArityNode *b){
  B32 result = 0;
  
  if (a->kind == b->kind){
    switch (a->kind){
      case DS_ArityKind_Empty:
      case DS_ArityKind_Singleton: result = 1; break;
      
      case DS_ArityKind_Arrow:
      {
        if (ds_aritynode_equal(a->arrow.domain, b->arrow.domain) &&
            ds_aritynode_equal(a->arrow.codomain, b->arrow.codomain)){
          result = 1;
        }
      }break;
      
      case DS_ArityKind_Tuple:
      {
        if (a->tuple.part_count == b->tuple.part_count){
          result = 1;
          for (U64 i = 0; i < a->tuple.part_count; i += 1){
            DS_ArityNode *a_part = a->tuple.parts[i];
            B32 found_match = 0;
            for (U64 j = 0; j < b->tuple.part_count; j += 1){
              DS_ArityNode *b_part = b->tuple.parts[j];
              if (b_part->coefficient == a_part->coefficient &&
                  ds_aritynode_equal(a_part, b_part)){
                found_match = 1;
                break;
              }
            }
            if (!found_match){
              result = 0;
              break;
            }
          }
        }
      }break;
    }
  }
  
  return(result);
}

// helpers

MR4TH_SYMBOL DS_ArityNode*
ds_aritynode__part_copy(Arena *arena, DS_ArityNode *node){
  Assert(node->kind != DS_ArityKind_Tuple);
  DS_ArityNode *result = push_array(arena, DS_ArityNode, 1);
  result->kind = node->kind;
  if (result->kind == DS_ArityKind_Arrow){
    result->arrow.domain = node->arrow.domain;
    result->arrow.codomain = node->arrow.codomain;
  }
  return(result);
}

MR4TH_SYMBOL DS_ArityNode*
ds_aritynode__part_find_match(DS_ArityNode **parts_first, DS_ArityNode **parts_opl, DS_ArityNode *key){
  DS_ArityNode **part = parts_first;
  DS_ArityNode *match = 0;
  for (; part < parts_opl; part += 1){
    if (ds_aritynode_equal(*part, key)){
      match = *part;
      break;
    }
  }
  return(match);
}