@@ -0,0 +1,222 @@
1+{-# LANGUAGE GeneralizedNewtypeDeriving #-}
2+{-# LANGUAGE OverloadedStrings #-}
3+
4+module Database.Bloodhound.Common.Script where
5+
6+import Bloodhound.Import
7+
8+import qualified Data.HashMap.Strict as HM
9+
10+import Database.V5.Bloodhound.Types.Internal.Newtypes
11+
12+newtype ScriptFields =
13+ ScriptFields (HM.HashMap ScriptFieldName ScriptFieldValue)
14+ deriving (Eq, Show)
15+
16+type ScriptFieldName = Text
17+type ScriptFieldValue = Value
18+
19+data Script =
20+ Script { scriptLanguage :: Maybe ScriptLanguage
21+ , scriptInline :: Maybe ScriptInline
22+ , scriptStored :: Maybe ScriptId
23+ , scriptParams :: Maybe ScriptParams
24+ } deriving (Eq, Show)
25+
26+newtype ScriptLanguage =
27+ ScriptLanguage Text deriving (Eq, Show, FromJSON, ToJSON)
28+
29+newtype ScriptInline =
30+ ScriptInline Text deriving (Eq, Show, FromJSON, ToJSON)
31+
32+newtype ScriptId =
33+ ScriptId Text deriving (Eq, Show, FromJSON, ToJSON)
34+
35+newtype ScriptParams =
36+ ScriptParams (HM.HashMap ScriptParamName ScriptParamValue)
37+ deriving (Eq, Show)
38+
39+type ScriptParamName = Text
40+type ScriptParamValue = Value
41+
42+data BoostMode =
43+ BoostModeMultiply
44+ | BoostModeReplace
45+ | BoostModeSum
46+ | BoostModeAvg
47+ | BoostModeMax
48+ | BoostModeMin deriving (Eq, Show)
49+
50+data ScoreMode =
51+ ScoreModeMultiply
52+ | ScoreModeSum
53+ | ScoreModeAvg
54+ | ScoreModeFirst
55+ | ScoreModeMax
56+ | ScoreModeMin deriving (Eq, Show)
57+
58+data FunctionScoreFunction =
59+ FunctionScoreFunctionScript Script
60+ | FunctionScoreFunctionRandom Seed
61+ | FunctionScoreFunctionFieldValueFactor FieldValueFactor
62+ deriving (Eq, Show)
63+
64+newtype Weight =
65+ Weight Float deriving (Eq, Show, FromJSON, ToJSON)
66+
67+newtype Seed =
68+ Seed Float deriving (Eq, Show, FromJSON, ToJSON)
69+
70+data FieldValueFactor =
71+ FieldValueFactor { fieldValueFactorField :: FieldName
72+ , fieldValueFactor :: Maybe Factor
73+ , fieldValueFactorModifier :: Maybe FactorModifier
74+ , fieldValueFactorMissing :: Maybe FactorMissingFieldValue
75+ } deriving (Eq, Show)
76+
77+newtype Factor =
78+ Factor Float deriving (Eq, Show, FromJSON, ToJSON)
79+
80+data FactorModifier =
81+ FactorModifierNone
82+ | FactorModifierLog
83+ | FactorModifierLog1p
84+ | FactorModifierLog2p
85+ | FactorModifierLn
86+ | FactorModifierLn1p
87+ | FactorModifierLn2p
88+ | FactorModifierSquare
89+ | FactorModifierSqrt
90+ | FactorModifierReciprocal deriving (Eq, Show)
91+
92+newtype FactorMissingFieldValue =
93+ FactorMissingFieldValue Float deriving (Eq, Show, FromJSON, ToJSON)
94+
95+instance ToJSON BoostMode where
96+ toJSON BoostModeMultiply = "multiply"
97+ toJSON BoostModeReplace = "replace"
98+ toJSON BoostModeSum = "sum"
99+ toJSON BoostModeAvg = "avg"
100+ toJSON BoostModeMax = "max"
101+ toJSON BoostModeMin = "min"
102+
103+instance FromJSON BoostMode where
104+ parseJSON = withText "BoostMode" parse
105+ where parse "multiply" = pure BoostModeMultiply
106+ parse "replace" = pure BoostModeReplace
107+ parse "sum" = pure BoostModeSum
108+ parse "avg" = pure BoostModeAvg
109+ parse "max" = pure BoostModeMax
110+ parse "min" = pure BoostModeMin
111+ parse bm = fail ("Unexpected BoostMode: " <> show bm)
112+
113+instance ToJSON ScoreMode where
114+ toJSON ScoreModeMultiply = "multiply"
115+ toJSON ScoreModeSum = "sum"
116+ toJSON ScoreModeFirst = "first"
117+ toJSON ScoreModeAvg = "avg"
118+ toJSON ScoreModeMax = "max"
119+ toJSON ScoreModeMin = "min"
120+
121+instance FromJSON ScoreMode where
122+ parseJSON = withText "ScoreMode" parse
123+ where parse "multiply" = pure ScoreModeMultiply
124+ parse "sum" = pure ScoreModeSum
125+ parse "first" = pure ScoreModeFirst
126+ parse "avg" = pure ScoreModeAvg
127+ parse "max" = pure ScoreModeMax
128+ parse "min" = pure ScoreModeMin
129+ parse sm = fail ("Unexpected ScoreMode: " <> show sm)
130+
131+functionScoreFunctionPair :: FunctionScoreFunction -> (Text, Value)
132+functionScoreFunctionPair (FunctionScoreFunctionScript functionScoreScript) =
133+ ("script_score", toJSON functionScoreScript)
134+
135+functionScoreFunctionPair (FunctionScoreFunctionRandom seed) =
136+ ("random_score", omitNulls [ "seed" .= seed ])
137+
138+functionScoreFunctionPair (FunctionScoreFunctionFieldValueFactor fieldValueFactor) =
139+ ("field_value_factor", toJSON fieldValueFactor)
140+
141+parseFunctionScoreFunction :: Object -> Parser FunctionScoreFunction
142+parseFunctionScoreFunction o =
143+ singleScript `taggedWith` "script_score"
144+ <|> singleRandom `taggedWith` "random_score"
145+ <|> singleFieldValueFactor `taggedWith` "field_value_factor"
146+ where taggedWith parser k = parser =<< o .: k
147+ singleScript = pure . FunctionScoreFunctionScript
148+ singleRandom o' = FunctionScoreFunctionRandom <$> o' .: "seed"
149+ singleFieldValueFactor = pure . FunctionScoreFunctionFieldValueFactor
150+
151+instance ToJSON ScriptFields where
152+ toJSON (ScriptFields x) = Object x
153+
154+instance FromJSON ScriptFields where
155+ parseJSON (Object o) = pure (ScriptFields o)
156+ parseJSON _ = fail "error parsing ScriptFields"
157+
158+instance ToJSON Script where
159+ toJSON (Script lang inline stored params) =
160+ object [ "script" .= omitNulls base ]
161+ where base = [ "lang" .= lang
162+ , "inline" .= inline
163+ , "stored" .= stored
164+ , "params" .= params ]
165+
166+instance FromJSON Script where
167+ parseJSON = withObject "Script" parse
168+ where parse o = o .: "script" >>= \o' ->
169+ Script
170+ <$> o' .:? "lang"
171+ <*> o' .:? "inline"
172+ <*> o' .:? "stored"
173+ <*> o' .:? "params"
174+
175+instance ToJSON ScriptParams where
176+ toJSON (ScriptParams x) = Object x
177+
178+instance FromJSON ScriptParams where
179+ parseJSON (Object o) = pure (ScriptParams o)
180+ parseJSON _ = fail "error parsing ScriptParams"
181+
182+instance ToJSON FieldValueFactor where
183+ toJSON (FieldValueFactor field factor modifier missing) =
184+ omitNulls base
185+ where base = [ "field" .= field
186+ , "factor" .= factor
187+ , "modifier" .= modifier
188+ , "missing" .= missing ]
189+
190+instance FromJSON FieldValueFactor where
191+ parseJSON = withObject "FieldValueFactor" parse
192+ where parse o = FieldValueFactor
193+ <$> o .: "field"
194+ <*> o .:? "factor"
195+ <*> o .:? "modifier"
196+ <*> o .:? "missing"
197+
198+instance ToJSON FactorModifier where
199+ toJSON FactorModifierNone = "none"
200+ toJSON FactorModifierLog = "log"
201+ toJSON FactorModifierLog1p = "log1p"
202+ toJSON FactorModifierLog2p = "log2p"
203+ toJSON FactorModifierLn = "ln"
204+ toJSON FactorModifierLn1p = "ln1p"
205+ toJSON FactorModifierLn2p = "ln2p"
206+ toJSON FactorModifierSquare = "square"
207+ toJSON FactorModifierSqrt = "sqrt"
208+ toJSON FactorModifierReciprocal = "reciprocal"
209+
210+instance FromJSON FactorModifier where
211+ parseJSON = withText "FactorModifier" parse
212+ where parse "none" = pure FactorModifierNone
213+ parse "log" = pure FactorModifierLog
214+ parse "log1p" = pure FactorModifierLog1p
215+ parse "log2p" = pure FactorModifierLog2p
216+ parse "ln" = pure FactorModifierLn
217+ parse "ln1p" = pure FactorModifierLn1p
218+ parse "ln2p" = pure FactorModifierLn2p
219+ parse "square" = pure FactorModifierSquare
220+ parse "sqrt" = pure FactorModifierSqrt
221+ parse "reciprocal" = pure FactorModifierReciprocal
222+ parse fm = fail ("Unexpected FactorModifier: " <> show fm)