prog_synthesis/02.hs

import Data.List (elemIndex, sort, find)
import Data.Maybe (isJust)

infixl 4 :+:

data Expr = Zero
          | Succ Expr
          | Len Expr
          | FirstZero Expr
          | Sort Expr
          | SubList Expr Expr Expr
          | Expr :+: Expr
          | Recursive Expr
          | ZeroList
          | InList
    deriving (Read, Show, Eq)

data Type = IntT | ListT
    deriving (Read, Show, Eq)
data Value = IntV Int | ListV [Int]
    deriving (Read, Show, Eq)

typeOf :: Expr -> Type
typeOf Zero = IntT
typeOf (Succ {}) = IntT
typeOf (Len {}) = IntT
typeOf (FirstZero {}) = IntT
typeOf _ = ListT

data Env = Env {input :: [Int], prog :: Expr, steps :: Int}

stepInEnv :: Env -> Env
stepInEnv env@(Env {input, prog, steps}) = env { steps = steps - 1 }

execInt :: Env -> Expr -> Maybe Int
execInt (Env {input, prog, steps=0}) _ = Nothing
execInt _ Zero = Just 0
execInt env (Succ expr) = (+) 1 <$> execInt (stepInEnv env) expr
execInt env (Len listExpr) = length <$> execList (stepInEnv env) listExpr
execInt env (FirstZero listExpr) = do value <- execList (stepInEnv env) listExpr
                                      0 `elemIndex` value
execInt _ _ = Nothing

execList :: Env -> Expr -> Maybe [Int]
execList (Env {input, prog, steps=0}) _ = Nothing
execList env (Sort listExpr) = sort <$> execList (stepInEnv env) listExpr
execList env (SubList listExpr exprFrom exprTo) = do valFrom  <- execInt (stepInEnv env) exprFrom
                                                     valTo <- execInt (stepInEnv env) exprTo
                                                     listValue <- execList (stepInEnv env) listExpr
                                                     return $ drop valFrom $ take (valTo + 1) listValue
execList env (exprLeft :+: exprRight) = do valLeft <- execList (stepInEnv env) exprLeft
                                           valRight <- execList (stepInEnv env) exprRight
                                           return $ valLeft ++ valRight
execList env (Recursive listExpr) = do listValue <- execList (stepInEnv env) listExpr
                                       if null listValue then Just [] else execList (stepInEnv $ env {input = listValue}) (prog env)
execList _ ZeroList = Just [0]
execList (Env {input, prog, steps}) InList = Just input
execList _ _ = Nothing

-- TODO: union
execProg :: [Int] -> Expr -> Maybe [Int]
execProg input expr = execList (Env {input, prog=expr, steps=20}) expr

terminals :: [Expr]
terminals = [Zero, ZeroList, InList] -- ,
             -- Succ Zero, Len InList]

concatShuffle :: [[Expr]] -> [Expr]
concatShuffle xxs = let xxs' = filter (not . null) xxs in
                    if null xxs' then [] else
                    map head xxs' ++ concatShuffle (map tail xxs')

nextSimpleExprsLists :: [Expr] -> [[Expr]]
nextSimpleExprsLists exprs = [[Succ e | e <- exprs, typeOf e == IntT],
                              [Sort e | e <- exprs, typeOf e == ListT],
                              [Recursive e | e <- exprs, typeOf e == ListT],
                              [FirstZero e | e <- exprs, typeOf e == ListT],
                              [Len e | e <- exprs, typeOf e == ListT]]

nextExprsLists :: [Expr] -> [[Expr]]
nextExprsLists exprs =  nextSimpleExprsLists exprs ++
                       [[e :+: e' | e <- exprs, typeOf e == ListT,
                                    e' <- exprs, typeOf e' == ListT],
                        [SubList e from to | e <- exprs, typeOf e == ListT,
                                             from <- exprs, typeOf from == IntT,
                                             to <- exprs, typeOf to == IntT]]

nextSimpleExprs :: [Expr] -> [Expr]
nextSimpleExprs exprs = (++) exprs $ concatShuffle $ nextSimpleExprsLists exprs

nextExprs :: [Expr] -> [Expr]
nextExprs exprs = (++) exprs $  concatShuffle $ nextExprsLists exprs

-- NOTE: slower version due to additional elem checks
-- nextExprsLists' :: [Expr] -> [Expr] -> [[Expr]]
-- nextExprsLists' prevExprs allExprs =  nextSimpleExprsLists prevExprs ++
--                         [[e :+: e' | e <- allExprs, typeOf e == ListT,
--                                      e' <- allExprs, typeOf e' == ListT,
--                                      e `elem` prevExprs || e' `elem` prevExprs],
--                          [SubList e from to | e <- allExprs, typeOf e == ListT,
--                                               from <- allExprs, typeOf from == IntT,
--                                               to <- allExprs, typeOf to == IntT,
--                                               e `elem` prevExprs || from `elem` prevExprs || to `elem` prevExprs]]

nextSimpleExprs' :: [Expr] -> [Expr]
nextSimpleExprs' = concatShuffle . nextSimpleExprsLists

-- TODO: check formula for three args
nextExprsLists' :: [Expr] -> [Expr] -> [[Expr]]
nextExprsLists' prevExprs allExprs =  nextSimpleExprsLists prevExprs ++
                        [[e :+: e' | e <- prevExprs, typeOf e == ListT,
                                     e' <- allExprs, typeOf e' == ListT],
                         [e :+: e' | e <- allExprs, typeOf e == ListT, e `notElem` prevExprs,
                                     e' <- prevExprs, typeOf e' == ListT],
                         [SubList e from to | e <- prevExprs, typeOf e == ListT,
                                              from <- allExprs, typeOf from == IntT,
                                              to <- allExprs, typeOf to == IntT],
                         [SubList e from to | e <- allExprs, typeOf e == ListT, e `notElem` prevExprs,
                                              from <- prevExprs, typeOf from == IntT,
                                              to <- allExprs, typeOf to == IntT],
                         [SubList e from to | e <- allExprs, typeOf e == ListT, e `notElem` prevExprs,
                                              from <- allExprs, typeOf from == IntT, from `notElem` prevExprs,
                                              to <- prevExprs, typeOf to == IntT]]

nextExprs' :: [Expr] -> [Expr] -> [Expr]
nextExprs' prevExprs allExprs = concatShuffle $ nextExprsLists'  prevExprs allExprs

data Example = Example {exampleInput :: [Int], exampleOutput :: [Int]}

-- check expr on all examples
isCorrect :: [Example] -> Expr -> Bool
isCorrect examples expr = all (\Example {exampleInput, exampleOutput} -> execProg exampleInput expr == Just exampleOutput) examples

isValid :: [Example] -> Expr -> Bool
isValid examples expr | typeOf expr == IntT = True
                      | otherwise =  all (\Example {exampleInput, exampleOutput} -> isJust $ execProg exampleInput expr) examples

-- check are exprs produce same results on all the examples
areSame :: [Example] -> Expr -> Expr -> Bool
areSame examples exprLeft exprRight | typeOf exprLeft == IntT = False
                                    | typeOf exprRight == IntT = False
                                    | otherwise = all (\Example {exampleInput, exampleOutput} -> let Just resLeft = execProg exampleInput exprLeft in
                                                                                                 let Just resRight = execProg exampleInput exprRight in
                                                                                                 resLeft /= [] && resLeft == resRight) examples

-----

upSyntesisStep :: [Example] -> [Expr] -> Either [Expr] Expr
upSyntesisStep examples exprs =  case find (isCorrect examples) exprs of
                                   Just answer -> Right answer
                                   Nothing -> let exprs' = filter (isValid examples) exprs in -- exclude invalid fragments
                                              let exprs'' = foldl (\acc expr -> if any (areSame examples expr) acc then acc else expr : acc) [] exprs' in -- merge same values
                                              Left $
                                                nextSimpleExprs $
                                                nextExprs exprs''

upSyntesisRec :: [Example] -> Int -> [Expr]  -> Maybe Expr
upSyntesisRec _ 0 _ = Nothing
upSyntesisRec examples steps exprs = case upSyntesisStep examples exprs of
                                       Right answer -> Just answer
                                       Left exprs' -> upSyntesisRec examples (steps - 1) exprs'


upSyntesis :: [Example] -> Int -> Maybe Expr
upSyntesis examples steps = upSyntesisRec examples steps $ nextSimpleExprs terminals

-----

upSyntesisStep' :: [Example] -> [Expr] -> [Expr] -> Either [Expr] Expr
upSyntesisStep' examples prevExprs allExprs =
                                case find (isCorrect examples) prevExprs of
                                  Just answer -> Right answer
                                  Nothing -> let allExprs' = filter (isValid examples) allExprs in -- exclude invalid fragments
                                             let prevExprs' = filter (isValid examples) prevExprs in -- exclude invalid fragments
                                             let allExprs'' = foldl (\acc expr -> if any (areSame examples expr) acc then acc else expr : acc) [] allExprs' in -- merge same values
                                             let prevExprs'' = foldl (\acc expr -> if any (areSame examples expr) acc then acc else expr : acc) [] prevExprs' in -- merge same values
                                             Left $
                                                    -- nextSimpleExprs $
                                                    nextExprs' prevExprs'' allExprs''

upSyntesisRec' :: [Example] -> Int -> [Expr] -> [Expr]  -> Maybe Expr
upSyntesisRec' _ 0 _ _ = Nothing
upSyntesisRec' examples steps prevExprs allExprs = case upSyntesisStep' examples prevExprs allExprs of
                                                    Right answer -> Just answer
                                                    Left exprs ->  upSyntesisRec' examples (steps - 1) exprs (allExprs ++ exprs)

upSyntesis' :: [Example] -> Int -> Maybe Expr
upSyntesis' examples steps = let terminals' = nextSimpleExprs terminals in
                             upSyntesisRec' examples steps terminals' terminals'

-----

exampleOf :: [Int] -> [Int] -> Example
exampleOf input output =  Example {exampleInput = input, exampleOutput = output}

sameExamplesExpected = InList
sameExamples = [exampleOf [1] [1], exampleOf [1,2] [1,2], exampleOf [1,2,3] [1,2,3], exampleOf [1,2,3,4] [1,2,3,4]]
revExamplesExpected = Recursive (SubList InList (Succ Zero) (Len InList)) :+: SubList InList Zero Zero
revExamples = [exampleOf [1] [1], exampleOf [1,2] [2,1], exampleOf [1,2,3] [3,2,1], exampleOf [1,2,3,4] [4,3,2,1]]

main = print $ upSyntesis' revExamples 4