{-# LANGUAGE PatternGuards, LambdaCase #-}

module Development.NSIS.Optimise(optimise) where

import Development.NSIS.Type
import Data.Generics.Uniplate.Data
import Data.List
import Data.Maybe


-- before: secret = 1021, primes = 109

optimise :: [NSIS] -> [NSIS]
optimise :: [NSIS] -> [NSIS]
optimise =
    -- allow Label 0
    ([NSIS] -> [NSIS]) -> [NSIS] -> [NSIS]
rep ([NSIS] -> [NSIS]
elimDeadLabel ([NSIS] -> [NSIS]) -> ([NSIS] -> [NSIS]) -> [NSIS] -> [NSIS]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [NSIS] -> [NSIS]
useLabel0) ([NSIS] -> [NSIS]) -> ([NSIS] -> [NSIS]) -> [NSIS] -> [NSIS]
forall b c a. (b -> c) -> (a -> b) -> a -> c
.
    -- disallow Label 0
    ([NSIS] -> [NSIS]) -> [NSIS] -> [NSIS]
rep ([NSIS] -> [NSIS]
elimDeadLabel ([NSIS] -> [NSIS]) -> ([NSIS] -> [NSIS]) -> [NSIS] -> [NSIS]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [NSIS] -> [NSIS]
elimAfterGoto ([NSIS] -> [NSIS]) -> ([NSIS] -> [NSIS]) -> [NSIS] -> [NSIS]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [NSIS] -> [NSIS]
deadAssign ([NSIS] -> [NSIS]) -> ([NSIS] -> [NSIS]) -> [NSIS] -> [NSIS]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [NSIS] -> [NSIS]
assignSwitch ([NSIS] -> [NSIS]) -> ([NSIS] -> [NSIS]) -> [NSIS] -> [NSIS]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [NSIS] -> [NSIS]
dullGoto ([NSIS] -> [NSIS]) -> ([NSIS] -> [NSIS]) -> [NSIS] -> [NSIS]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [NSIS] -> [NSIS]
knownCompare ([NSIS] -> [NSIS]) -> ([NSIS] -> [NSIS]) -> [NSIS] -> [NSIS]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [NSIS] -> [NSIS]
elimLabeledGoto ([NSIS] -> [NSIS]) -> ([NSIS] -> [NSIS]) -> [NSIS] -> [NSIS]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [NSIS] -> [NSIS]
elimDeadVar)


rep :: ([NSIS] -> [NSIS]) -> [NSIS] -> [NSIS]
rep :: ([NSIS] -> [NSIS]) -> [NSIS] -> [NSIS]
rep [NSIS] -> [NSIS]
f [NSIS]
x = Int -> [NSIS] -> [NSIS]
g ([NSIS] -> Int
forall {from}. Data from => from -> Int
measure [NSIS]
x) [NSIS]
x
    where
        g :: Int -> [NSIS] -> [NSIS]
g Int
n1 [NSIS]
x1 = if Int
n2 Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
n1 then Int -> [NSIS] -> [NSIS]
g Int
n2 [NSIS]
x2 else [NSIS]
x2
            where x2 :: [NSIS]
x2 = [NSIS] -> [NSIS]
f ([NSIS] -> [NSIS]) -> [NSIS] -> [NSIS]
forall a b. (a -> b) -> a -> b
$ [NSIS] -> [NSIS]
f ([NSIS] -> [NSIS]) -> [NSIS] -> [NSIS]
forall a b. (a -> b) -> a -> b
$ [NSIS] -> [NSIS]
f ([NSIS] -> [NSIS]) -> [NSIS] -> [NSIS]
forall a b. (a -> b) -> a -> b
$ [NSIS] -> [NSIS]
f [NSIS]
x1
                  n2 :: Int
n2 = [NSIS] -> Int
forall {from}. Data from => from -> Int
measure [NSIS]
x2
        measure :: from -> Int
measure from
x = [NSIS] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length (from -> [NSIS]
forall from to. Biplate from to => from -> [to]
universeBi from
x :: [NSIS])


useLabel0 :: [NSIS] -> [NSIS]
useLabel0 :: [NSIS] -> [NSIS]
useLabel0 = (NSIS -> NSIS) -> [NSIS] -> [NSIS]
forall a b. (a -> b) -> [a] -> [b]
map (([NSIS] -> [NSIS]) -> NSIS -> NSIS
forall from to. Biplate from to => (to -> to) -> from -> from
descendBi [NSIS] -> [NSIS]
useLabel0) ([NSIS] -> [NSIS]) -> ([NSIS] -> [NSIS]) -> [NSIS] -> [NSIS]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [NSIS] -> [NSIS]
f
    where
        f :: [NSIS] -> [NSIS]
f (NSIS
x:Labeled Label
next:[NSIS]
xs)
            | [NSIS] -> Bool
forall a. [a] -> Bool
forall (t :: * -> *) a. Foldable t => t a -> Bool
null (NSIS -> [NSIS]
forall on. Uniplate on => on -> [on]
children NSIS
x :: [NSIS]) -- must not be a block with nested instructions
            = (Label -> Label) -> NSIS -> NSIS
forall from to. Biplate from to => (to -> to) -> from -> from
descendBi (\Label
i -> if Label
i Label -> Label -> Bool
forall a. Eq a => a -> a -> Bool
== Label
next then Int -> Label
Label Int
0 else Label
i) NSIS
x NSIS -> [NSIS] -> [NSIS]
forall a. a -> [a] -> [a]
: Label -> NSIS
Labeled Label
next NSIS -> [NSIS] -> [NSIS]
forall a. a -> [a] -> [a]
: [NSIS] -> [NSIS]
f [NSIS]
xs
        f (NSIS
x:[NSIS]
xs) = NSIS
x NSIS -> [NSIS] -> [NSIS]
forall a. a -> [a] -> [a]
: [NSIS] -> [NSIS]
f [NSIS]
xs
        f [] = []


-- Label whose next statement is a good,
elimLabeledGoto :: [NSIS] -> [NSIS]
elimLabeledGoto :: [NSIS] -> [NSIS]
elimLabeledGoto [NSIS]
x = (NSIS -> NSIS) -> [NSIS] -> [NSIS]
forall from to. Biplate from to => (to -> to) -> from -> from
transformBi NSIS -> NSIS
f [NSIS]
x
    where
        f :: NSIS -> NSIS
f (Labeled Label
x) = Label -> NSIS
Labeled Label
x
        f NSIS
x | [NSIS] -> Bool
forall a. [a] -> Bool
forall (t :: * -> *) a. Foldable t => t a -> Bool
null (NSIS -> [NSIS]
forall on. Uniplate on => on -> [on]
children NSIS
x) = (Label -> Label) -> NSIS -> NSIS
forall from to. Biplate from to => (to -> to) -> from -> from
descendBi Label -> Label
moveBounce NSIS
x
            | Bool
otherwise = NSIS
x

        moveBounce :: Label -> Label
moveBounce Label
x = Label -> Maybe Label -> Label
forall a. a -> Maybe a -> a
fromMaybe Label
x (Maybe Label -> Label) -> Maybe Label -> Label
forall a b. (a -> b) -> a -> b
$ Label -> [(Label, Label)] -> Maybe Label
forall a b. Eq a => a -> [(a, b)] -> Maybe b
lookup Label
x [(Label, Label)]
bounce
        bounce :: [(Label, Label)]
bounce = (([NSIS] -> [(Label, Label)]) -> [[NSIS]] -> [(Label, Label)])
-> [[NSIS]] -> ([NSIS] -> [(Label, Label)]) -> [(Label, Label)]
forall a b c. (a -> b -> c) -> b -> a -> c
flip ([NSIS] -> [(Label, Label)]) -> [[NSIS]] -> [(Label, Label)]
forall (t :: * -> *) a b. Foldable t => (a -> [b]) -> t a -> [b]
concatMap ([NSIS] -> [[NSIS]]
forall on. Uniplate on => on -> [on]
universe [NSIS]
x) (([NSIS] -> [(Label, Label)]) -> [(Label, Label)])
-> ([NSIS] -> [(Label, Label)]) -> [(Label, Label)]
forall a b. (a -> b) -> a -> b
$ \case
            Labeled Label
x:Goto Label
y:[NSIS]
_ -> [(Label
x,Label
y)]
            Labeled Label
x:Labeled Label
y:[NSIS]
_ -> [(Label
x,Label
y)]
            [NSIS]
_ -> []


-- Delete variables which are only assigned, never read from
elimDeadVar :: [NSIS] -> [NSIS]
elimDeadVar :: [NSIS] -> [NSIS]
elimDeadVar [NSIS]
x = ([NSIS] -> [NSIS]) -> [NSIS] -> [NSIS]
forall on. Uniplate on => (on -> on) -> on -> on
transform [NSIS] -> [NSIS]
f [NSIS]
x
    where
        f :: [NSIS] -> [NSIS]
f (Assign Var
x Val
_:[NSIS]
xs) | Var
x Var -> [Var] -> Bool
forall a. Eq a => a -> [a] -> Bool
forall (t :: * -> *) a. (Foldable t, Eq a) => a -> t a -> Bool
`elem` [Var]
unused = [NSIS]
xs
        f [NSIS]
xs = [NSIS]
xs

        unused :: [Var]
unused = [Var] -> [Var]
forall a. Eq a => [a] -> [a]
nub [Var]
assign [Var] -> [Var] -> [Var]
forall a. Eq a => [a] -> [a] -> [a]
\\ [Var] -> [Var]
forall a. Eq a => [a] -> [a]
nub [Var]
used
        used :: [Var]
used = [Var]
every [Var] -> [Var] -> [Var]
forall a. Eq a => [a] -> [a] -> [a]
\\ [Var]
assign
        every :: [Var]
every = [NSIS] -> [Var]
forall from to. Biplate from to => from -> [to]
universeBi [NSIS]
x
        assign :: [Var]
assign = [Var
x | Assign Var
x Val
_ <- [NSIS] -> [NSIS]
forall from to. Biplate from to => from -> [to]
universeBi [NSIS]
x]

jumpy :: NSIS -> Bool
jumpy Goto{} = Bool
True
jumpy StrCmpS{} = Bool
True
jumpy IntCmp{} = Bool
True
jumpy IfErrors{} = Bool
True
jumpy IfFileExists{} = Bool
True
jumpy MessageBox{} = Bool
True
jumpy NSIS
_ = Bool
False


-- Eliminate any code after a goto, until a label
elimAfterGoto :: [NSIS] -> [NSIS]
elimAfterGoto :: [NSIS] -> [NSIS]
elimAfterGoto [NSIS]
x = ([NSIS] -> [NSIS]) -> [NSIS] -> [NSIS]
forall from to. Biplate from to => (to -> to) -> from -> from
transformBi [NSIS] -> [NSIS]
f [NSIS]
x
    where
        f :: [NSIS] -> [NSIS]
f (NSIS
x:[NSIS]
xs) | NSIS -> Bool
jumpy NSIS
x = NSIS
x NSIS -> [NSIS] -> [NSIS]
forall a. a -> [a] -> [a]
: [NSIS] -> [NSIS]
g [NSIS]
xs
        f [NSIS]
x = [NSIS]
x

        g :: [NSIS] -> [NSIS]
g (Labeled Label
x:[NSIS]
xs) = Label -> NSIS
Labeled Label
xNSIS -> [NSIS] -> [NSIS]
forall a. a -> [a] -> [a]
:[NSIS]
xs
        g (NSIS
x:[NSIS]
xs) = [NSIS] -> [NSIS]
g [NSIS]
xs
        g [NSIS]
x = [NSIS]
x


-- Be careful to neither introduce or remove label based errors
elimDeadLabel :: [NSIS] -> [NSIS]
elimDeadLabel :: [NSIS] -> [NSIS]
elimDeadLabel [NSIS]
x = ([NSIS] -> [NSIS]) -> [NSIS] -> [NSIS]
forall on. Uniplate on => (on -> on) -> on -> on
transform [NSIS] -> [NSIS]
f [NSIS]
x
    where
        f :: [NSIS] -> [NSIS]
f (Labeled Label
x:[NSIS]
xs) | Label
x Label -> [Label] -> Bool
forall a. Eq a => a -> [a] -> Bool
forall (t :: * -> *) a. (Foldable t, Eq a) => a -> t a -> Bool
`elem` [Label]
unused = [NSIS]
xs
        f [NSIS]
xs = [NSIS]
xs

        unused :: [Label]
unused = [Label] -> [Label]
forall a. Eq a => [a] -> [a]
nub [Label]
label [Label] -> [Label] -> [Label]
forall a. Eq a => [a] -> [a] -> [a]
\\ [Label] -> [Label]
forall a. Eq a => [a] -> [a]
nub [Label]
gotos
        gotos :: [Label]
gotos = [Label]
every [Label] -> [Label] -> [Label]
forall a. Eq a => [a] -> [a] -> [a]
\\ [Label]
label
        every :: [Label]
every = [NSIS] -> [Label]
forall from to. Biplate from to => from -> [to]
universeBi [NSIS]
x
        label :: [Label]
label = [Label
x | Labeled Label
x <- [NSIS] -> [NSIS]
forall from to. Biplate from to => from -> [to]
universeBi [NSIS]
x]


dullGoto :: [NSIS] -> [NSIS]
dullGoto :: [NSIS] -> [NSIS]
dullGoto = ([NSIS] -> [NSIS]) -> [NSIS] -> [NSIS]
forall on. Uniplate on => (on -> on) -> on -> on
transform [NSIS] -> [NSIS]
f
    where
        f :: [NSIS] -> [NSIS]
f (Goto Label
l1:Labeled Label
l2:[NSIS]
xs) | Label
l1 Label -> Label -> Bool
forall a. Eq a => a -> a -> Bool
== Label
l2 = Label -> NSIS
Labeled Label
l2 NSIS -> [NSIS] -> [NSIS]
forall a. a -> [a] -> [a]
: [NSIS]
xs
        f [NSIS]
x = [NSIS]
x


-- A tricky one! Comparison after jump
knownCompare :: [NSIS] -> [NSIS]
knownCompare :: [NSIS] -> [NSIS]
knownCompare [NSIS]
x = ([NSIS] -> [NSIS]) -> [NSIS] -> [NSIS]
forall on. Uniplate on => (on -> on) -> on -> on
transform [NSIS] -> [NSIS]
f [NSIS]
x
    where
        f :: [NSIS] -> [NSIS]
f (Assign Var
var Val
val : StrCmpS Val
a Val
b Label
yes Label
no : [NSIS]
xs)
            | Val
a Val -> Val -> Bool
forall a. Eq a => a -> a -> Bool
== [Var -> Val_
Var_ Var
var], Just Bool
eq <- Val -> Val -> Maybe Bool
isEqual Val
b Val
val
            = Var -> Val -> NSIS
Assign Var
var Val
val NSIS -> [NSIS] -> [NSIS]
forall a. a -> [a] -> [a]
: Label -> NSIS
Goto (if Bool
eq then Label
yes else Label
no) NSIS -> [NSIS] -> [NSIS]
forall a. a -> [a] -> [a]
: [NSIS]
xs

        -- grows, but only a finite amount
        f (Assign Var
var Val
val : Labeled Label
l : StrCmpS Val
a Val
b Label
yes Label
no : [NSIS]
xs)
            | Val
a Val -> Val -> Bool
forall a. Eq a => a -> a -> Bool
== [Var -> Val_
Var_ Var
var], Just Bool
eq <- Val -> Val -> Maybe Bool
isEqual Val
b Val
val
            = Var -> Val -> NSIS
Assign Var
var Val
val NSIS -> [NSIS] -> [NSIS]
forall a. a -> [a] -> [a]
: Label -> NSIS
Goto (if Bool
eq then Label
yes else Label
no) NSIS -> [NSIS] -> [NSIS]
forall a. a -> [a] -> [a]
: Label -> NSIS
Labeled Label
l NSIS -> [NSIS] -> [NSIS]
forall a. a -> [a] -> [a]
: Val -> Val -> Label -> Label -> NSIS
StrCmpS Val
a Val
b Label
yes Label
no NSIS -> [NSIS] -> [NSIS]
forall a. a -> [a] -> [a]
: [NSIS]
xs

        f (Assign Var
var Val
val : NSIS
c : [NSIS]
xs) | NSIS -> Bool
jumpy NSIS
c = Var -> Val -> NSIS
Assign Var
var Val
val NSIS -> [NSIS] -> [NSIS]
forall a. a -> [a] -> [a]
: (Label -> Label) -> NSIS -> NSIS
forall from to. Biplate from to => (to -> to) -> from -> from
transformBi Label -> Label
g NSIS
c NSIS -> [NSIS] -> [NSIS]
forall a. a -> [a] -> [a]
: [NSIS]
xs
            where
                g :: Label -> Label
g Label
l | Just (StrCmpS Val
a Val
b Label
yes Label
no) <- Label -> [(Label, NSIS)] -> Maybe NSIS
forall a b. Eq a => a -> [(a, b)] -> Maybe b
lookup Label
l [(Label, NSIS)]
cmps
                    , Val
a Val -> Val -> Bool
forall a. Eq a => a -> a -> Bool
== [Var -> Val_
Var_ Var
var], Just Bool
eq <- Val -> Val -> Maybe Bool
isEqual Val
b Val
val
                    = if Bool
eq then Label
yes else Label
no
                g Label
l = Label
l
        f [NSIS]
x = [NSIS]
x

        cmps :: [(Label, NSIS)]
cmps = [(Label
l,NSIS
cmp) | Labeled Label
l : cmp :: NSIS
cmp@StrCmpS{} : [NSIS]
_ <- [NSIS] -> [[NSIS]]
forall from to. Biplate from to => from -> [to]
universeBi [NSIS]
x]


isEqual :: Val -> Val -> Maybe Bool
isEqual :: Val -> Val -> Maybe Bool
isEqual Val
x Val
y | Val
x Val -> Val -> Bool
forall a. Eq a => a -> a -> Bool
== Val
y = Bool -> Maybe Bool
forall a. a -> Maybe a
Just Bool
True
            | Val -> Bool
isLit Val
x, Val -> Bool
isLit Val
y = Bool -> Maybe Bool
forall a. a -> Maybe a
Just Bool
False
            | Bool
otherwise = Maybe Bool
forall a. Maybe a
Nothing
    where
        isLit :: Val -> Bool
isLit = (Val_ -> Bool) -> Val -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
all Val_ -> Bool
isLiteral
        isLiteral :: Val_ -> Bool
isLiteral Literal{} = Bool
True
        isLiteral Val_
_ = Bool
False


assignSwitch :: [NSIS] -> [NSIS]
assignSwitch :: [NSIS] -> [NSIS]
assignSwitch = ([NSIS] -> [NSIS]) -> [NSIS] -> [NSIS]
forall on. Uniplate on => (on -> on) -> on -> on
transform [NSIS] -> [NSIS]
f
    where
        -- this rule just switches the assignment, back and forth, ad infinitum
        -- not very principled!
        f :: [NSIS] -> [NSIS]
f (IntOp Var
out1 Val
a String
b Val
c : Assign Var
other ([Var_ Var
out2]) : [NSIS]
xs)
            | Var
out1 Var -> Var -> Bool
forall a. Eq a => a -> a -> Bool
== Var
out2
            = Var -> Val -> String -> Val -> NSIS
IntOp Var
other Val
a String
b Val
c NSIS -> [NSIS] -> [NSIS]
forall a. a -> [a] -> [a]
: Var -> Val -> NSIS
Assign Var
out1 ([Var -> Val_
Var_ Var
other]) NSIS -> [NSIS] -> [NSIS]
forall a. a -> [a] -> [a]
: [NSIS]
xs
        f [NSIS]
x = [NSIS]
x


deadAssign :: [NSIS] -> [NSIS]
deadAssign :: [NSIS] -> [NSIS]
deadAssign = ([NSIS] -> [NSIS]) -> [NSIS] -> [NSIS]
forall on. Uniplate on => (on -> on) -> on -> on
transform [NSIS] -> [NSIS]
f
    where
        f :: [NSIS] -> [NSIS]
f (Assign Var
v Val
x:[NSIS]
xs) | Var -> [NSIS] -> Bool
isDead Var
v [NSIS]
xs = [NSIS]
xs
        f [NSIS]
xs = [NSIS]
xs

        isDead :: Var -> [NSIS] -> Bool
isDead Var
v (Labeled Label
_:[NSIS]
xs) = Var -> [NSIS] -> Bool
isDead Var
v [NSIS]
xs
        isDead Var
v (Assign Var
v2 Val
x:[NSIS]
xs) = Var
v Var -> [Var] -> Bool
forall (t :: * -> *) a. (Foldable t, Eq a) => a -> t a -> Bool
`notElem` Val -> [Var]
forall from to. Biplate from to => from -> [to]
universeBi Val
x Bool -> Bool -> Bool
&& (Var
v Var -> Var -> Bool
forall a. Eq a => a -> a -> Bool
== Var
v2 Bool -> Bool -> Bool
|| Var -> [NSIS] -> Bool
isDead Var
v [NSIS]
xs)
        isDead Var
v [NSIS]
_ = Bool
False