content: {-# OPTIONS_GHC -Wall #-}
  
  module Optimize.Case
    ( optimize,
    )
  where
  
  import AST.Canonical qualified as Can
  import AST.Optimized qualified as Opt
  import Control.Arrow (second)
  import Data.Map ((!))
  import Data.Map qualified as Map
  import Data.Maybe qualified as Maybe
  import Data.Name qualified as Name
  import Optimize.DecisionTree qualified as DT
  
  -- OPTIMIZE A CASE EXPRESSION
  
  optimize :: Name.Name -> Name.Name -> [(Can.Pattern, Opt.Expr)] -> Opt.Expr
  optimize temp root optBranches =
    let (patterns, indexedBranches) =
          unzip (zipWith indexify [0 ..] optBranches)
  
        decider = treeToDecider (DT.compile patterns)
        targetCounts = countTargets decider
  
        (choices, maybeJumps) =
          unzip (map (createChoices targetCounts) indexedBranches)
     in Opt.Case
          temp
          root
          (insertChoices (Map.fromList choices) decider)
          (Maybe.catMaybes maybeJumps)
  
  indexify :: Int -> (a, b) -> ((a, Int), (Int, b))
  indexify index (pattern, branch) =
    ( (pattern, index),
      (index, branch)
    )
  
  -- TREE TO DECIDER
  --
  -- Decision trees may have some redundancies, so we convert them to a Decider
  -- which has special constructs to avoid code duplication when possible.
  
  treeToDecider :: DT.DecisionTree -> Opt.Decider Int
  treeToDecider tree =
    case tree of
      DT.Match target ->
        Opt.Leaf target
      -- zero options
      DT.Decision _ [] Nothing ->
        error "compiler bug, somehow created an empty decision tree"
      -- one option
      DT.Decision _ [(_, subTree)] Nothing ->
        treeToDecider subTree
      DT.Decision _ [] (Just subTree) ->
        treeToDecider subTree
      -- two options
      DT.Decision path [(test, successTree)] (Just failureTree) ->
        toChain path test successTree failureTree
      DT.Decision path [(test, successTree), (_, failureTree)] Nothing ->
        toChain path test successTree failureTree
      -- many options
      DT.Decision path edges Nothing ->
        let (necessaryTests, fallback) =
              (init edges, snd (last edges))
         in Opt.FanOut
              path
              (map (second treeToDecider) necessaryTests)
              (treeToDecider fallback)
      DT.Decision path edges (Just fallback) ->
        Opt.FanOut path (map (second treeToDecider) edges) (treeToDecider fallback)
  
  toChain :: DT.Path -> DT.Test -> DT.DecisionTree -> DT.DecisionTree -> Opt.Decider Int
  toChain path test successTree failureTree =
    let failure =
          treeToDecider failureTree
     in case treeToDecider successTree of
          Opt.Chain testChain success subFailure
            | failure == subFailure ->
                Opt.Chain ((path, test) : testChain) success failure
          success ->
            Opt.Chain [(path, test)] success failure
  
  -- INSERT CHOICES
  --
  -- If a target appears exactly once in a Decider, the corresponding expression
  -- can be inlined. Whether things are inlined or jumps is called a "choice".
  
  countTargets :: Opt.Decider Int -> Map.Map Int Int
  countTargets decisionTree =
    case decisionTree of
      Opt.Leaf target ->
        Map.singleton target 1
      Opt.Chain _ success failure ->
        Map.unionWith (+) (countTargets success) (countTargets failure)
      Opt.FanOut _ tests fallback ->
        Map.unionsWith (+) (map countTargets (fallback : map snd tests))
  
  createChoices ::
    Map.Map Int Int ->
    (Int, Opt.Expr) ->
    ((Int, Opt.Choice), Maybe (Int, Opt.Expr))
  createChoices targetCounts (target, branch) =
    if targetCounts ! target == 1
      then
        ( (target, Opt.Inline branch),
          Nothing
        )
      else
        ( (target, Opt.Jump target),
          Just (target, branch)
        )
  
  insertChoices ::
    Map.Map Int Opt.Choice ->
    Opt.Decider Int ->
    Opt.Decider Opt.Choice
  insertChoices choiceDict decider =
    let go =
          insertChoices choiceDict
     in case decider of
          Opt.Leaf target ->
            Opt.Leaf (choiceDict ! target)
          Opt.Chain testChain success failure ->
            Opt.Chain testChain (go success) (go failure)
          Opt.FanOut path tests fallback ->
            Opt.FanOut path (map (second go) tests) (go fallback)
  