{-# LANGUAGE BangPatterns #-}

module Copilot.Zephyr.Main (zephyr) where

import Language.Copilot (Spec, interpret, reify)
import Copilot.Compile.C99 (compile)
import Copilot.Zephyr.Internals
import System.IO
import System.Directory
import System.IO.Temp (withSystemTempDirectory, createTempDirectory)
import System.FilePath
import System.Exit
import Data.List (isInfixOf, intercalate)
import Data.Maybe
import qualified Data.Map as M
import qualified Data.Set as S
import Options.Applicative

-- | Typically your program's main will be implemented using this.
-- For example:
--
-- > {-# LANGUAGE RebindableSyntax #-}
-- > 
-- > import Copilot.Zephyr.Board.Generic
-- > 
-- > main = zephyr $ do
-- >	led0 =: flashing
-- > 	delay =: MilliSeconds (constant 100)
--
-- Running this program compiles the `Sketch` into C code using copilot
-- and generates a Zephyr app in the directory "generated". That app
-- can be built and uploaded to your board using Zephyr, the same as any
-- other Zephyr app. See Zephyr's documentation for details.
--
-- This also supports interpreting a `Sketch`, without loading it onto a
-- board. Run the program with parameters "-i 4" to display what it
-- would do on the first 4 iterations. The output will look something like
-- this:
--
-- > gpio_pin_set_led0:  k_msleep:
-- > (false)             (100)
-- > (true)              (100)
-- > (false)             (100)
-- > (true)              (100)
zephyr :: Sketch () -> IO ()
zephyr :: Sketch () -> IO ()
zephyr Sketch ()
s = CmdLine -> IO ()
go forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< forall a. ParserInfo a -> IO a
execParser ParserInfo CmdLine
opts
  where
	opts :: ParserInfo CmdLine
opts = forall a. Parser a -> InfoMod a -> ParserInfo a
info (Parser CmdLine
parseCmdLine forall (f :: * -> *) a b. Applicative f => f a -> f (a -> b) -> f b
<**> forall a. Parser (a -> a)
helper)
		( forall a. InfoMod a
fullDesc
		forall a. Semigroup a => a -> a -> a
<> forall a. String -> InfoMod a
progDesc String
"Run this program with no options to generate an Zephyr program."
		)

	go :: CmdLine -> IO ()
go CmdLine
o = case (Maybe Spec
mspec, CmdLine -> Maybe Integer
interpretSteps CmdLine
o) of
		(Maybe Spec
Nothing, Maybe Integer
_) -> do
			Handle -> String -> IO ()
hPutStrLn Handle
stderr String
"This Sketch does not do anything."
			forall a. IO a
exitFailure
		(Just Spec
spec, Just Integer
n) -> Integer -> Spec -> IO ()
interpret Integer
n Spec
spec
		(Just Spec
spec, Maybe Integer
Nothing) -> do
			let name :: String
name = String
"generated"
			String -> Spec -> Framework -> IO ()
writeMain String
name Spec
spec Framework
f'
			String -> [KConfig] -> IO ()
writePrjConf String
name [KConfig]
kconfigs
			String -> DeviceTree -> IO ()
writeAppOverlay String
name DeviceTree
devicetree
			String -> IO ()
writeCMakeLists String
name
	
	(Maybe Spec
mspec, Framework
f) = forall ctx a.
Context ctx =>
GenSketch ctx a -> (Maybe Spec, GenFramework ctx)
evalSketch Sketch ()
s
	
	-- Strict evaluation because this may throw errors,
	-- and we want to throw them before starting compilation.
	!(Framework
f', [KConfig]
kconfigs, DeviceTree
devicetree) = Framework -> (Framework, [KConfig], DeviceTree)
finalizeFramework Framework
f

data CmdLine = CmdLine
	{ CmdLine -> Maybe Integer
interpretSteps :: Maybe Integer
	}


parseCmdLine :: Parser CmdLine
parseCmdLine :: Parser CmdLine
parseCmdLine = Maybe Integer -> CmdLine
CmdLine
	forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall (f :: * -> *) a. Alternative f => f a -> f (Maybe a)
optional (forall a. ReadM a -> Mod OptionFields a -> Parser a
option forall a. Read a => ReadM a
auto
		( forall (f :: * -> *) a. HasName f => String -> Mod f a
long String
"interpret"
		forall a. Semigroup a => a -> a -> a
<> forall (f :: * -> *) a. HasName f => Char -> Mod f a
short Char
'i'
		forall a. Semigroup a => a -> a -> a
<> forall (f :: * -> *) a. String -> Mod f a
help String
"use copilot to interpret the program, displaying what it would do"
		forall a. Semigroup a => a -> a -> a
<> forall (f :: * -> *) a. HasMetavar f => String -> Mod f a
metavar String
"NUM"
		))

writeMain :: String -> Spec -> Framework -> IO ()
writeMain :: String -> Spec -> Framework -> IO ()
writeMain String
name Spec
spec Framework
framework = do
	-- This could be a lot prettier, unfortunately copilot only exports
	-- an interface that writes the generated code to a file.
	-- And, the .c file includes a .h file that will make it fail to
	-- build, so that include has to be filtered out.
	forall (m :: * -> *) a.
(MonadIO m, MonadMask m) =>
String -> (String -> m a) -> m a
withSystemTempDirectory String
"copilot" forall a b. (a -> b) -> a -> b
$ \String
toptmpdir -> do
		String
mytmpdir <- String -> String -> IO String
createTempDirectory String
toptmpdir String
"copilot"
		forall a. Spec' a -> IO Spec
reify Spec
spec forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= String -> Spec -> IO ()
compile (String
mytmpdir String -> String -> String
</> String
"copilot")
		[String]
c <- String -> [String]
lines forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> String -> IO String
readFile (String
mytmpdir String -> String -> String
</> String
"copilot.c")
		let c' :: [String]
c' = forall a. (a -> Bool) -> [a] -> [a]
filter (Bool -> Bool
Prelude.not forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a. Eq a => [a] -> [a] -> Bool
isInfixOf String
"#include \"") [String]
c
		Bool -> String -> IO ()
createDirectoryIfMissing Bool
False String
name
		String -> String -> IO ()
writeFile (String
name String -> String -> String
</> String
"main.c") forall a b. (a -> b) -> a -> b
$ 
			[String] -> String
unlines forall a b. (a -> b) -> a -> b
$ forall a b. (a -> b) -> [a] -> [b]
map CLine -> String
fromCLine forall a b. (a -> b) -> a -> b
$
				Framework -> [CLine] -> [CLine]
sketchFramework Framework
framework (forall a b. (a -> b) -> [a] -> [b]
map String -> CLine
CLine [String]
c')

writeCMakeLists :: String -> IO ()
writeCMakeLists :: String -> IO ()
writeCMakeLists String
name = do
	Bool -> String -> IO ()
createDirectoryIfMissing Bool
False String
name
	String -> String -> IO ()
writeFile (String
name String -> String -> String
</> String
"CMakeLists.txt") forall a b. (a -> b) -> a -> b
$ [String] -> String
unlines
		[ String
"cmake_minimum_required(VERSION 3.20.0)"
		, String
"find_package(Zephyr REQUIRED HINTS $ENV{ZEPHYR_BASE})"
		, String
"project(" forall a. [a] -> [a] -> [a]
++ String
name forall a. [a] -> [a] -> [a]
++ String
")"
		, String
""
		, String
"target_sources(app PRIVATE main.c)"
		]

data KConfig = KConfig String

writePrjConf :: String -> [KConfig] -> IO ()
writePrjConf :: String -> [KConfig] -> IO ()
writePrjConf String
name [KConfig]
kconfigs = do
	Bool -> String -> IO ()
createDirectoryIfMissing Bool
False String
name
	String -> String -> IO ()
writeFile (String
name String -> String -> String
</> String
"prj.conf") forall a b. (a -> b) -> a -> b
$ [String] -> String
unlines forall a b. (a -> b) -> a -> b
$
		forall a b. (a -> b) -> [a] -> [b]
map (\(KConfig String
s) -> String
s) [KConfig]
kconfigs

newtype DeviceTree = DeviceTree String

writeAppOverlay :: String -> DeviceTree -> IO ()
writeAppOverlay :: String -> DeviceTree -> IO ()
writeAppOverlay String
name (DeviceTree String
s) = do
	Bool -> String -> IO ()
createDirectoryIfMissing Bool
False String
name
	String -> String -> IO ()
writeFile (String
name String -> String -> String
</> String
"app.overlay") String
s

-- Makes an Zephyr C program, using a Framework, and a list of lines of C
-- code generated by Copilot.
sketchFramework :: Framework -> [CLine] -> [CLine]
sketchFramework :: Framework -> [CLine] -> [CLine]
sketchFramework Framework
f [CLine]
ccode = forall a b. (a -> b) -> [a] -> [b]
map String -> CLine
CLine forall a b. (a -> b) -> a -> b
$ forall (t :: * -> *) a. Foldable t => t [a] -> [a]
concat
	[
		[ String
"/* automatically generated, do not edit */"
		, String
blank
		, String
"#include <zephyr.h>"
		, String
blank
		]
	, forall a b. (a -> b) -> [a] -> [b]
map CLine -> String
fromCLine ([CChunk] -> [CLine]
fromchunks (forall ctx. GenFramework ctx -> [CChunk]
defines Framework
f))
	, [String
blank]
	, forall a b. (a -> b) -> [a] -> [b]
map CLine -> String
fromCLine [CLine]
ccode
	, [String
blank]
	,
		[ String
"void main(void)"
		]
	, [String] -> [String]
codeblock forall a b. (a -> b) -> a -> b
$ forall (t :: * -> *) a. Foldable t => t [a] -> [a]
concat
		[ forall a b. (a -> b) -> [a] -> [b]
map CLine -> String
fromCLine ([CChunk] -> [CLine]
fromchunks (forall ctx. GenFramework ctx -> [CChunk]
earlySetups Framework
f))
		, forall a b. (a -> b) -> [a] -> [b]
map CLine -> String
fromCLine ([CChunk] -> [CLine]
fromchunks (forall ctx. GenFramework ctx -> [CChunk]
setups Framework
f))
		, [String
"while (1)"]
		, [String] -> [String]
codeblock forall a b. (a -> b) -> a -> b
$ forall (t :: * -> *) a. Foldable t => t [a] -> [a]
concat
			[ forall a b. (a -> b) -> [a] -> [b]
map CLine -> String
fromCLine ([CChunk] -> [CLine]
fromchunks (forall ctx. GenFramework ctx -> [CChunk]
loops Framework
f))
			, [ String
"step();" ]
			]
		]
	]
  where
	blank :: String
blank = String
""
	indent :: String -> String
indent String
l = String
"  " forall a. Semigroup a => a -> a -> a
<> String
l
	codeblock :: [String] -> [String]
codeblock [String]
l = [String
"{"] forall a. Semigroup a => a -> a -> a
<> forall a b. (a -> b) -> [a] -> [b]
map String -> String
indent [String]
l forall a. Semigroup a => a -> a -> a
<> [String
"}"]
	-- If two CChunks are identical, only include it once.
	-- This can happen when eg, the same setup code is generated
	-- to use a resource that's accessed more than once in a program.
	-- Note: Does not preserve order of chunks in the list.
	fromchunks :: [CChunk] -> [CLine]
	fromchunks :: [CChunk] -> [CLine]
fromchunks [CChunk]
cl = forall (t :: * -> *) a b. Foldable t => (a -> [b]) -> t a -> [b]
concatMap (\(CChunk [CLine]
l) -> [CLine]
l) forall a b. (a -> b) -> a -> b
$
		forall a. Set a -> [a]
S.toList forall a b. (a -> b) -> a -> b
$ forall a. Ord a => [a] -> Set a
S.fromList [CChunk]
cl

finalizeFramework :: Framework -> (Framework, [KConfig], DeviceTree)
finalizeFramework :: Framework -> (Framework, [KConfig], DeviceTree)
finalizeFramework Framework
f = 
	let ([[CChunk]]
pindefines, [[CChunk]]
pinsetups) = forall a b. [(a, b)] -> ([a], [b])
unzip forall a b. (a -> b) -> a -> b
$ forall a b. (a -> Maybe b) -> [a] -> [b]
mapMaybe (Zephyr, Set PinMode) -> Maybe ([CChunk], [CChunk])
setuppin
	    	(forall k a. Map k a -> [(k, a)]
M.toList (forall ctx. GenFramework ctx -> Map ctx (Set PinMode)
pinmodes Framework
f))
	    includes :: [CChunk]
includes = if Bool -> Bool
not (forall (t :: * -> *) a. Foldable t => t a -> Bool
null (forall (t :: * -> *) a. Foldable t => t [a] -> [a]
concat [[CChunk]]
pindefines))
		then [CLine] -> [CChunk]
mkCChunk 
			[ String -> CLine
CLine String
"#include <device.h>"
			, String -> CLine
CLine String
"#include <devicetree.h>"
			, String -> CLine
CLine String
"#include <drivers/gpio.h>"
			]
		else []
	    f' :: Framework
f' = Framework
f
		{ defines :: [CChunk]
defines = [CChunk]
includes forall a. Semigroup a => a -> a -> a
<> forall (t :: * -> *) a. Foldable t => t [a] -> [a]
concat [[CChunk]]
pindefines forall a. Semigroup a => a -> a -> a
<> forall ctx. GenFramework ctx -> [CChunk]
defines Framework
f
		, setups :: [CChunk]
setups = forall (t :: * -> *) a. Foldable t => t [a] -> [a]
concat [[CChunk]]
pinsetups forall a. Semigroup a => a -> a -> a
<> forall ctx. GenFramework ctx -> [CChunk]
setups Framework
f
		}
	    kconfigs :: [KConfig]
kconfigs = if Bool -> Bool
not (forall (t :: * -> *) a. Foldable t => t a -> Bool
null (forall (t :: * -> *) a. Foldable t => t [a] -> [a]
concat [[CChunk]]
pindefines))
		then [String -> KConfig
KConfig String
"CONFIG_GPIO=y"]
		else []
	    devicetree :: DeviceTree
devicetree = String -> DeviceTree
DeviceTree forall a b. (a -> b) -> a -> b
$ [String] -> String
unlines forall a b. (a -> b) -> a -> b
$
	    	forall (t :: * -> *) a b. Foldable t => (a -> [b]) -> t a -> [b]
concatMap Zephyr -> [String]
mkdevicetree (forall k a. Map k a -> [k]
M.keys (forall ctx. GenFramework ctx -> Map ctx (Set PinMode)
pinmodes Framework
f))
	   in (Framework
f', [KConfig]
kconfigs, DeviceTree
devicetree)
  where
	setuppin :: (Zephyr, Set PinMode) -> Maybe ([CChunk], [CChunk])
setuppin (Zephyr (GPIOAlias String
n) GPIOAddress
_, Set PinMode
s)
		| Set PinMode
s forall a. Eq a => a -> a -> Bool
== forall a. a -> Set a
S.singleton PinMode
OutputMode = forall a. a -> Maybe a
Just forall a b. (a -> b) -> a -> b
$
			( String -> [CChunk]
definepin String
n
			, String -> String -> [CChunk]
setmode String
n ([String] -> String
ored [String
"GPIO_OUTPUT", String -> String
gpioflags String
n])
			)
		| Set PinMode
s forall a. Eq a => a -> a -> Bool
== forall a. a -> Set a
S.singleton PinMode
InputMode = forall a. a -> Maybe a
Just forall a b. (a -> b) -> a -> b
$
			( String -> [CChunk]
definepin String
n
			, String -> String -> [CChunk]
setmode String
n ([String] -> String
ored [String
"GPIO_INPUT", String -> String
gpioflags String
n])
			)
		| Set PinMode
s forall a. Eq a => a -> a -> Bool
== forall a. a -> Set a
S.singleton PinMode
InputPullupMode = forall a. a -> Maybe a
Just forall a b. (a -> b) -> a -> b
$
			( String -> [CChunk]
definepin String
n
			, String -> String -> [CChunk]
setmode String
n ([String] -> String
ored [String
"GPIO_PULL_UP", String -> String
gpioflags String
n])
			)
		| Set PinMode
s forall a. Eq a => a -> a -> Bool
== forall a. Ord a => [a] -> Set a
S.fromList [PinMode
InputMode, PinMode
InputPullupMode] = forall a. a -> Maybe a
Just forall a b. (a -> b) -> a -> b
$
			( String -> [CChunk]
definepin String
n
			, String -> String -> [CChunk]
setmode String
n ([String] -> String
ored [String
"GPIO_INPUT", String
"GPIO_PULL_UP", String -> String
gpioflags String
n])
			)
		| forall a. Set a -> Bool
S.null Set PinMode
s = forall a. Maybe a
Nothing
		| Bool
otherwise = forall a. HasCallStack => String -> a
error forall a b. (a -> b) -> a -> b
$
			String
"The program uses pin " forall a. [a] -> [a] -> [a]
++ forall a. Show a => a -> String
show String
n forall a. [a] -> [a] -> [a]
++
			String
" in multiple ways in different places (" forall a. [a] -> [a] -> [a]
++
			[String] -> String
unwords (forall a b. (a -> b) -> [a] -> [b]
map forall a. Show a => a -> String
show (forall a. Set a -> [a]
S.toList Set PinMode
s)) forall a. [a] -> [a] -> [a]
++ String
"). " forall a. [a] -> [a] -> [a]
++
			String
"This is not currently supported by zephyr-copilot."
	
	definepin :: String -> [CChunk]
definepin String
n = [CLine] -> [CChunk]
mkCChunk
		[ String -> CLine
CLine forall a b. (a -> b) -> a -> b
$ String
"#define " forall a. Semigroup a => a -> a -> a
<> String -> String
gpionode String
n forall a. Semigroup a => a -> a -> a
<> String
" DT_ALIAS(" forall a. Semigroup a => a -> a -> a
<> String
n forall a. Semigroup a => a -> a -> a
<> String
")"
		, String -> CLine
CLine forall a b. (a -> b) -> a -> b
$ String
"#if DT_NODE_HAS_STATUS(" forall a. Semigroup a => a -> a -> a
<> String -> String
gpionode String
n forall a. Semigroup a => a -> a -> a
<> String
", okay)"
		, String -> CLine
CLine forall a b. (a -> b) -> a -> b
$ String
"#define " forall a. Semigroup a => a -> a -> a
<> String -> String
gpiolabel String
n 
			forall a. Semigroup a => a -> a -> a
<> String
" DT_GPIO_LABEL(" forall a. Semigroup a => a -> a -> a
<> String -> String
gpionode String
n forall a. Semigroup a => a -> a -> a
<> String
", gpios)"
		, String -> CLine
CLine forall a b. (a -> b) -> a -> b
$ String
"#define " forall a. Semigroup a => a -> a -> a
<> String -> String
gpiopin String
n 
			forall a. Semigroup a => a -> a -> a
<> String
" DT_GPIO_PIN(" forall a. Semigroup a => a -> a -> a
<> String -> String
gpionode String
n forall a. Semigroup a => a -> a -> a
<> String
", gpios)"
		, String -> CLine
CLine forall a b. (a -> b) -> a -> b
$ String
"#define " forall a. Semigroup a => a -> a -> a
<> String -> String
gpioflags String
n 
			forall a. Semigroup a => a -> a -> a
<> String
" DT_GPIO_FLAGS(" forall a. Semigroup a => a -> a -> a
<> String -> String
gpionode String
n forall a. Semigroup a => a -> a -> a
<> String
", gpios)"
		, String -> CLine
CLine forall a b. (a -> b) -> a -> b
$ String
"#else"
		, String -> CLine
CLine forall a b. (a -> b) -> a -> b
$ String
"#error \"Unsupported board: devicetree alias is not" 
			forall a. Semigroup a => a -> a -> a
<> String
" defined for " forall a. Semigroup a => a -> a -> a
<> String
n forall a. Semigroup a => a -> a -> a
<> String
"\""
		, String -> CLine
CLine forall a b. (a -> b) -> a -> b
$ String
"#define " forall a. Semigroup a => a -> a -> a
<> String -> String
gpiolabel String
n forall a. Semigroup a => a -> a -> a
<> String
""
		, String -> CLine
CLine forall a b. (a -> b) -> a -> b
$ String
"#define " forall a. Semigroup a => a -> a -> a
<> String -> String
gpiopin String
n forall a. Semigroup a => a -> a -> a
<> String
" 0"
		, String -> CLine
CLine forall a b. (a -> b) -> a -> b
$ String
"#define " forall a. Semigroup a => a -> a -> a
<> String -> String
gpioflags String
n forall a. Semigroup a => a -> a -> a
<> String
" 0"
		, String -> CLine
CLine forall a b. (a -> b) -> a -> b
$ String
"#endif"
		, String -> CLine
CLine forall a b. (a -> b) -> a -> b
$ String
"const struct device *" forall a. Semigroup a => a -> a -> a
<> String -> String
pinDevVar String
n forall a. Semigroup a => a -> a -> a
<> String
";"
		]

	setmode :: String -> String -> [CChunk]
setmode String
n String
v = [CLine] -> [CChunk]
mkCChunk
		[ String -> CLine
CLine forall a b. (a -> b) -> a -> b
$ String -> String
pinDevVar String
n forall a. Semigroup a => a -> a -> a
<> String
" = device_get_binding(" 
			forall a. Semigroup a => a -> a -> a
<> String -> String
gpiolabel String
n forall a. Semigroup a => a -> a -> a
<> String
");"
		, String -> CLine
CLine forall a b. (a -> b) -> a -> b
$ String
"gpio_pin_configure(" 
			forall a. Semigroup a => a -> a -> a
<> String -> String
pinDevVar String
n forall a. Semigroup a => a -> a -> a
<> String
", "
			forall a. Semigroup a => a -> a -> a
<> String -> String
gpiopin String
n forall a. Semigroup a => a -> a -> a
<> String
", " 
			forall a. Semigroup a => a -> a -> a
<> String
v forall a. Semigroup a => a -> a -> a
<> String
");" 
		]
	
	gpiopin :: String -> String
gpiopin String
n = String -> String
pinDevDef String
n
	gpiolabel :: String -> String
gpiolabel String
n = String -> String
pinDevDef (String
n forall a. Semigroup a => a -> a -> a
<> String
"_LABEL")
	gpioflags :: String -> String
gpioflags String
n = String -> String
pinDevDef (String
n forall a. Semigroup a => a -> a -> a
<> String
"_FLAGS")
	gpionode :: String -> String
gpionode String
n = String -> String
pinDevNode String
n
	
	ored :: [String] -> String
ored = forall a. [a] -> [[a]] -> [a]
intercalate String
" | "
	
	mkdevicetree :: Zephyr -> [String]
mkdevicetree (Zephyr GPIOAlias
_ GPIOAddress
GPIOAddressBuiltIn) = []
	mkdevicetree (Zephyr (GPIOAlias String
n) (GPIOAddress String
addr)) =
		[ String
"/ {"
		, String
"\tmypins {"
		, String
"\t\tcompatible = \"gpio-keys\";"
		, String
"\t\t" forall a. [a] -> [a] -> [a]
++ String
n forall a. [a] -> [a] -> [a]
++ String
": " forall a. [a] -> [a] -> [a]
++ String
n forall a. [a] -> [a] -> [a]
++ String
"{"
		, String
"\t\t\tgpios = <&" forall a. [a] -> [a] -> [a]
++ String
addr forall a. [a] -> [a] -> [a]
++ String
" 0>;"
		, String
"\t\t\tlabel = \"" forall a. [a] -> [a] -> [a]
++ String
n forall a. [a] -> [a] -> [a]
++ String
"\";"
		, String
"\t\t};"
		, String
"\t};"
		, String
"\taliases {"
		, String
"\t\t" forall a. [a] -> [a] -> [a]
++ forall a b. (a -> b) -> [a] -> [b]
map Char -> Char
toalias String
n forall a. [a] -> [a] -> [a]
++ String
" = &" forall a. [a] -> [a] -> [a]
++ String
n forall a. [a] -> [a] -> [a]
++ String
";"
		, String
"\t};"
		, String
"};"
		]
	  where
		toalias :: Char -> Char
toalias Char
'_' = Char
'-'
		toalias Char
c = Char
c