<commit_before>{-# LANGUAGE LambdaCase, TupleSections #-}
module Transformations.Optimising.SparseCaseOptimisation where

import qualified Data.Map as Map
import Data.Functor.Foldable as Foldable
import Grin
import TypeEnv

sparseCaseOptimisation :: (TypeEnv, Exp) -> (TypeEnv, Exp)
sparseCaseOptimisation (typeEnv, exp) = (typeEnv, ana builder exp) where
  builder :: Exp -> ExpF Exp
  builder = \case
    ECase val@(Var name) alts ->
      ECaseF val
        [ alt
        | alt@(Alt cpat body) <- alts
        , possible varType cpat
        ] where varType = variableType typeEnv name

    exp -> project exp

  possible :: Type -> CPat -> Bool
  possible (T_NodeSet nodeSet) cpat = case cpat of
    NodePat tag _args -> Map.member tag nodeSet
    _ -> False

  possible ty@T_SimpleType{} cpat = case cpat of
    LitPat lit -> ty == typeOfLit lit
    _ -> False

  possible _ _ = True -- bypass everything else
<commit_msg>fix default pattern handling in sparse case elimination<commit_after>{-# LANGUAGE LambdaCase, TupleSections #-}
module Transformations.Optimising.SparseCaseOptimisation where

import qualified Data.Map as Map
import Data.Set (Set)
import qualified Data.Set as Set
import Data.Functor.Foldable as Foldable
import Grin
import TypeEnv

sparseCaseOptimisation :: (TypeEnv, Exp) -> (TypeEnv, Exp)
sparseCaseOptimisation (typeEnv, exp) = (typeEnv, ana builder exp) where
  builder :: Exp -> ExpF Exp
  builder = \case
    ECase val@(Var name) alts ->
      ECaseF val
        [ alt
        | alt@(Alt cpat body) <- alts
        , possible varType allPatTags cpat
        ] where varType     = variableType typeEnv name :: Type
                allPatTags  = Set.fromList [tag | Alt (NodePat tag _) _ <- alts]

    exp -> project exp

  possible :: Type -> Set Tag -> CPat -> Bool
  possible (T_NodeSet nodeSet) allPatTags cpat = case cpat of
    NodePat tag _args -> Map.member tag nodeSet
    -- HINT: the default case is redundant if normal cases fully cover the domain
    DefaultPat -> Set.isProperSubsetOf allPatTags (Map.keysSet nodeSet)
    _ -> False

  possible ty@T_SimpleType{} _ cpat = case cpat of
    LitPat lit -> ty == typeOfLit lit
    DefaultPat -> True -- HINT: the value domain is unknown, it is not possible to prove if it overlaps or it is fully covered
    _ -> False

  possible _ _ _ = True -- bypass everything else

{'fraction_non_alphanumeric': 0.0829805249788315, 'fraction_numerical': 0.0, 'mean_word_length': 3.039316239316239, 'pattern_counts': {'":': 0, '<': 6, '<?xml version=': 0, '>': 26, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '4891343', 'n_tokens_mistral': 739, 'n_tokens_neox': 709, 'n_words': 259}