Fix function argument editing

ui/src/ifql/components/FuncArg.tsx

@@ -15,6 +15,7 @@ interface Props {
   value: string | boolean
   type: string
   bodyID: string
+  declarationID: string
   onChangeArg: OnChangeArg
   onGenerateScript: () => void
 }
@@ -26,10 +27,11 @@ class FuncArg extends PureComponent<Props> {
       argKey,
       value,
       type,
-      funcName,
-      funcID,
-      onChangeArg,
       bodyID,
+      funcID,
+      funcName,
+      onChangeArg,
+      declarationID,
       onGenerateScript,
     } = this.props
 
@@ -62,6 +64,7 @@ class FuncArg extends PureComponent<Props> {
             funcID={funcID}
             bodyID={bodyID}
             onChangeArg={onChangeArg}
+            declarationID={declarationID}
             onGenerateScript={onGenerateScript}
           />
         )
@@ -72,9 +75,9 @@ class FuncArg extends PureComponent<Props> {
           <FuncArgBool
             value={this.boolValue}
             argKey={argKey}
+            bodyID={bodyID}
             funcID={funcID}
             onChangeArg={onChangeArg}
-            bodyID={bodyID}
             onGenerateScript={onGenerateScript}
           />
         )

ui/src/ifql/components/FuncArgInput.tsx

@@ -8,6 +8,7 @@ interface Props {
   value: string
   type: string
   bodyID: string
+  declarationID: string
   onChangeArg: OnChangeArg
   onGenerateScript: () => void
 }
@@ -44,12 +45,13 @@ class FuncArgInput extends PureComponent<Props> {
   }
 
   private handleChange = (e: ChangeEvent<HTMLInputElement>) => {
-    const {funcID, argKey, bodyID} = this.props
+    const {funcID, argKey, bodyID, declarationID} = this.props
 
     this.props.onChangeArg({
       funcID,
       key: argKey,
       value: e.target.value,
+      declarationID,
       bodyID,
     })
   }

ui/src/ifql/components/FuncArgs.tsx

@@ -8,13 +8,20 @@ interface Props {
   func: Func
   bodyID: string
   onChangeArg: OnChangeArg
+  declarationID: string
   onGenerateScript: () => void
 }
 
 @ErrorHandling
 export default class FuncArgs extends PureComponent<Props> {
   public render() {
-    const {bodyID, func, onChangeArg, onGenerateScript} = this.props
+    const {
+      func,
+      bodyID,
+      onChangeArg,
+      declarationID,
+      onGenerateScript,
+    } = this.props
 
     return (
       <div className="func-args">
@@ -29,6 +36,7 @@ export default class FuncArgs extends PureComponent<Props> {
               funcID={func.id}
               funcName={func.name}
               onChangeArg={onChangeArg}
+              declarationID={declarationID}
               onGenerateScript={onGenerateScript}
             />
           )

ui/src/ifql/components/FuncNode.tsx

@@ -30,7 +30,13 @@ export default class FuncNode extends PureComponent<Props, State> {
   }
 
   public render() {
-    const {bodyID, func, onChangeArg, onGenerateScript} = this.props
+    const {
+      func,
+      bodyID,
+      onChangeArg,
+      declarationID,
+      onGenerateScript,
+    } = this.props
     const {isOpen} = this.state
 
     return (
@@ -43,6 +49,7 @@ export default class FuncNode extends PureComponent<Props, State> {
             func={func}
             bodyID={bodyID}
             onChangeArg={onChangeArg}
+            declarationID={declarationID}
             onGenerateScript={onGenerateScript}
           />
         )}

ui/src/ifql/containers/IFQLPage.tsx

@@ -1,11 +1,12 @@
 import React, {PureComponent} from 'react'
 
 import {connect} from 'react-redux'
+import _ from 'lodash'
 
 import TimeMachine from 'src/ifql/components/TimeMachine'
 import KeyboardShortcuts from 'src/shared/components/KeyboardShortcuts'
 import {Suggestion, FlatBody} from 'src/types/ifql'
-import {InputArg, Handlers, DeleteFuncNodeArgs} from 'src/types/ifql'
+import {InputArg, Handlers, DeleteFuncNodeArgs, Func} from 'src/types/ifql'
 
 import {bodyNodes} from 'src/ifql/helpers'
 import {getSuggestions, getAST} from 'src/ifql/apis'
@@ -108,7 +109,7 @@ export class IFQLPage extends PureComponent<Props, State> {
   }
 
   private handleGenerateScript = (): void => {
-    this.getASTResponse(this.expressionsToScript)
+    this.getASTResponse(this.bodyToScript)
   }
 
   private handleChangeArg = ({
@@ -116,30 +117,41 @@ export class IFQLPage extends PureComponent<Props, State> {
     value,
     generate,
     funcID,
+    declarationID = '',
     bodyID,
   }: InputArg): void => {
-    const body = this.state.body.map(expression => {
-      if (expression.id !== bodyID) {
-        return expression
+    const body = this.state.body.map(b => {
+      if (b.id !== bodyID) {
+        return b
       }
 
-      const funcs = expression.funcs.map(f => {
-        if (f.id !== funcID) {
-          return f
-        }
-
-        const args = f.args.map(a => {
-          if (a.key === key) {
-            return {...a, value}
+      if (declarationID) {
+        const declarations = b.declarations.map(d => {
+          if (d.id !== declarationID) {
+            return d
           }
 
-          return a
+          const functions = this.editFuncArgs({
+            funcs: d.funcs,
+            funcID,
+            key,
+            value,
+          })
+
+          return {...d, funcs: functions}
         })
 
-        return {...f, args}
+        return {...b, declarations}
+      }
+
+      const funcs = this.editFuncArgs({
+        funcs: b.funcs,
+        funcID,
+        key,
+        value,
       })
 
-      return {...expression, funcs}
+      return {...b, funcs}
     })
 
     this.setState({body}, () => {
@@ -149,9 +161,32 @@ export class IFQLPage extends PureComponent<Props, State> {
     })
   }
 
-  private get expressionsToScript(): string {
-    return this.state.body.reduce((acc, expression) => {
-      return `${acc + this.funcsToScript(expression.funcs)}\n\n`
+  private editFuncArgs = ({funcs, funcID, key, value}): Func[] => {
+    return funcs.map(f => {
+      if (f.id !== funcID) {
+        return f
+      }
+
+      const args = f.args.map(a => {
+        if (a.key === key) {
+          return {...a, value}
+        }
+
+        return a
+      })
+
+      return {...f, args}
+    })
+  }
+
+  private get bodyToScript(): string {
+    return this.state.body.reduce((acc, b) => {
+      if (b.declarations.length) {
+        const funcs = _.get(b, 'declarations.0.funcs', [])
+        return `${acc}${this.funcsToScript(funcs)}\n\n`
+      }
+
+      return `${acc}${this.funcsToScript(b.funcs)}\n\n`
     }, '')
   }
 
@@ -184,13 +219,13 @@ export class IFQLPage extends PureComponent<Props, State> {
   }
 
   private handleAddNode = (name: string, bodyID: string): void => {
-    const script = this.state.body.reduce((acc, expression) => {
-      if (expression.id === bodyID) {
-        const {funcs} = expression
+    const script = this.state.body.reduce((acc, body) => {
+      if (body.id === bodyID) {
+        const {funcs} = body
         return `${acc}${this.funcsToScript(funcs)}\n\t|> ${name}()\n\n`
       }
 
-      return acc + expression.source
+      return acc + body.source
     }, '')
 
     this.getASTResponse(script)

ui/src/types/ifql.ts

@@ -24,6 +24,7 @@ export interface DeleteFuncNodeArgs {
 export interface InputArg {
   funcID: string
   bodyID: string
+  declarationID?: string
   key: string
   value: string | boolean
   generate?: boolean

ui/test/ifql/components/FuncArg.test.tsx

@@ -7,6 +7,7 @@ const setup = () => {
     funcID: '',
     bodyID: '',
     funcName: '',
+    declarationID: '',
     argKey: '',
     value: '',
     type: '',